diff --git a/sotodlib/coords/det_match.py b/sotodlib/coords/det_match.py index ea9bc7806..aa4075e9e 100644 --- a/sotodlib/coords/det_match.py +++ b/sotodlib/coords/det_match.py @@ -65,23 +65,26 @@ class PointingConfig: ----- fp_file: str Path to focal-plane file that is used by the optics module. + rx_file: str + Path to optics tube YAML file that is used by the optics module. wafer_slot: str Wafer slot of the UFM. For example: "ws0" - tel_type: str - Tel type for the optics model. Either "SAT" or "LAT" + platform: str + Which platform to use for the optics model. zemax_path: str - If running for a "LAT" tel_type, the path to the zemax file must be specified. + If running for the lat platform, the path to the zemax file must be specified. roll: float Rotation about the line of sight. For the LAT this is elev - 60 - corotator. For the SAT this is -1*boresight. tube_slot: str/int - If running for a "LAT" tel_type, the tube slot must be specified. + If running for the lat platform, the tube slot must be specified. Either the tube name as a string or the tube number as an int. """ fp_file: str + rx_file: str wafer_slot: str - tel_type: str + platform: str zemax_path: Optional[str] = None roll: Optional[float] = 0 tube_slot: Optional[Union[str, int]] = None @@ -92,6 +95,10 @@ class PointingConfig: fp_pars: dict = field(init=False) def __post_init__(self): + if self.platform.lower() in ["satp1", "satp2", "satp3"]: + self.tel_type = "SAT" + elif self.platform.lower() == "lat": + self.tel_type = "LAT" if self.tel_type == 'LAT': if self.zemax_path is None: raise ValueError("zemax path must be set for 'LAT' tel_type") @@ -104,22 +111,37 @@ def __post_init__(self): self.fp_pars = optics.get_ufm_to_fp_pars( self.tel_type, self.wafer_slot, self.fp_file ) - self.dx = self.fp_pars['dx'] - self.dy = self.fp_pars['dy'] - self.theta = np.deg2rad(self.fp_pars['theta']) - def get_pointing(self, x, y, pol=0): - xp = x * np.cos(self.theta) - y * np.sin(self.theta) + self.dx - yp = x * np.sin(self.theta) + y * np.cos(self.theta) + self.dy + self.rx_pars = optics.get_fp_to_rx_pars( + self.tube_slot, self.rx_file) + def get_pointing(self, x, y, pol=0): if self.tel_type.upper() == 'SAT': - xi, eta, gamma = optics.SAT_focal_plane( - None, x=xp, y=yp, pol=pol, roll=self.roll + xi, eta, gamma = optics.get_focal_plane( + None, + x=x, + y=y, + pol=pol, + roll=self.roll, + telescope_flavor="SAT", + tube_slot=self.tube_slot, + wafer_slot=self.wafer_slot, + ufm_to_fp_pars=self.fp_pars, + fp_to_rx_pars=self.rx_pars, ) elif self.tel_type.upper() == 'LAT': - xi, eta, gamma = optics.LAT_focal_plane( - None, self.zemax_path, x=xp, y=yp, pol=pol, - roll=self.roll, tube_slot=self.tube_slot + xi, eta, gamma = optics.get_focal_plane( + None, + x=x, + y=y, + pol=pol, + roll=self.roll, + telescope_flavor="LAT", + tube_slot=self.tube_slot, + wafer_slot=self.wafer_slot, + ufm_to_fp_pars=self.fp_pars, + fp_to_rx_pars=self.rx_pars, + zemax_path=self.zemax_path, ) return xi, eta, gamma @@ -140,6 +162,8 @@ class Resonator: smurf_subband: int = -1 readout_id: str = '' + wafer_slot: str = 'ws.' + xi: float = np.nan eta: float = np.nan gamma: float = np.nan @@ -194,7 +218,7 @@ def apply_design_properties(smurf_res, design_res, in_place=False, apply_pointin 'det_pol', 'det_freq', 'det_bandpass', 'det_angle_raw_deg', 'det_angle_actual_deg', 'det_type', 'det_id', 'is_optical', 'mux_bondpad', 'mux_subband', 'mux_band', 'mux_channel', - 'mux_layout_pos' + 'mux_layout_pos', 'wafer_slot' ] # for LF @@ -285,7 +309,8 @@ def as_array(self): return np.array(data, dtype=dtype) @classmethod - def from_aman(cls, aman, stream_id, det_cal=None, name=None, pointing: Optional[AxisManager]=None): + def from_aman(cls, aman, stream_id, det_cal=None, name=None, pointing: Optional[AxisManager]=None, + ignore_north_south=False): """ Load a resonator set from a Context object based on an obs_id @@ -315,7 +340,10 @@ def from_aman(cls, aman, stream_id, det_cal=None, name=None, pointing: Optional[ resonators = [] for i, ri in enumerate(np.where(m)[0]): band, channel = aman.det_info.smurf.band[ri], aman.det_info.smurf.channel[ri] - is_north = north_is_highband ^ (band < 4) + if not ignore_north_south: + is_north = north_is_highband ^ (band < 4) + else: + is_north = 1 readout_id = aman.det_info.readout_id[ri] bg = det_cal.bg[ri] res_freq=aman.det_info.smurf.frequency[ri] @@ -336,7 +364,7 @@ def from_aman(cls, aman, stream_id, det_cal=None, name=None, pointing: Optional[ @classmethod def from_tunefile(cls, tunefile, name=None, north_is_highband=True, - resfit_file=None, bgmap_file=None): + resfit_file=None, bgmap_file=None, ignore_north_south=False): """ Creates an instance based on a smurf-tune file. If a resfit or bgmap file is included, that data will be added to the Resonance objects as @@ -363,7 +391,10 @@ def from_tunefile(cls, tunefile, name=None, north_is_highband=True, continue for res_idx, d in _v['resonances'].items(): - is_north = north_is_highband ^ (band < 4) + if not ignore_north_south: + is_north = north_is_highband ^ (band < 4) + else: + is_north = 1 res = Resonator( idx=idx, smurf_res_idx=res_idx, res_freq=d['freq'], smurf_band=band, is_north=is_north @@ -385,7 +416,8 @@ def from_tunefile(cls, tunefile, name=None, north_is_highband=True, @classmethod def from_wafer_info_file(cls, wafer_info_file, array_name, name=None, - pt_cfg: Optional[PointingConfig]=None): + pt_cfg: Optional[PointingConfig]=None, + ignore_north_south=False): """ Initialize a ResSet from a wafer info file. This is a file that contains detector design information. @@ -415,7 +447,10 @@ def from_wafer_info_file(cls, wafer_info_file, array_name, name=None, resonators = [] idx = 0 for r in wafer_array: - is_north = r['dets:wafer.coax'] == b'N' + if not ignore_north_south: + is_north = r['dets:wafer.coax'] == b'N' + else: + is_north = 1 kwargs = dict( idx=idx, det_id=r['dets:det_id'].decode(), @@ -457,7 +492,8 @@ def from_wafer_info_file(cls, wafer_info_file, array_name, name=None, @classmethod def from_solutions(cls, sol_file, north_is_highband=True, name=None, - fp_pars=None, platform='SAT', zemax_path=None): + fp_pars=None, platform='LAT', zemax_path=None, + ignore_north_south=False): """ Creates an instance from an input-solution file. This will include both design data, along with smurf-band and smurf-channel info. Resonance frequencies used here are the VNA @@ -475,7 +511,7 @@ def from_solutions(cls, sol_file, north_is_highband=True, name=None, Result of the function ``sotododlib.coords.optics.get_ufm_to_fp_pars``. If this is None, detector positions will not be mapped to pointing angles. platform (str): - 'SAT' or 'LAT'. Used to determine which focal plane function to + Which platform. Used to determine which focal plane function to use for pointing zemax_path (str): zemax path, required to get pointing for LAT optics @@ -505,7 +541,10 @@ def _int(val, null_val=None): continue # is_north = north_is_highband ^ (_int(d['smurf_band']) < 4) # is_north = d['is_north'].lower().strip() == 'true' - is_north = _int(d['bias_line']) < 6 + if not ignore_north_south: + is_north = _int(d['bias_line']) < 6 + else: + is_north = 1 is_optical = (d['is_optical'].lower() == 'true') kwargs = dict( @@ -745,7 +784,7 @@ def _get_biadjacency_matrix(self) -> np.ndarray: mat = np.zeros((len(self.src), len(self.dst)), dtype=float) - # N/S mismatch + # N/S mismatch (all N for LF) m = src_arr['is_north'][:, None] != dst_arr['is_north'][None, :] mat[m] = np.inf diff --git a/sotodlib/coords/det_match_solutions.py b/sotodlib/coords/det_match_solutions.py index 2399e04f4..b28061366 100644 --- a/sotodlib/coords/det_match_solutions.py +++ b/sotodlib/coords/det_match_solutions.py @@ -6,7 +6,6 @@ import os import numpy as np import yaml -from copy import deepcopy from scipy import interpolate from tqdm.auto import tqdm, trange from copy import deepcopy @@ -32,14 +31,16 @@ class SolutionsCfg: Directory where results should be stored. wafer_info_path: str Path to the wafer_info h5 file. - tel_type: str - Tel type for the optics model. Either "SAT" or "LAT" + platform: str + Telescope platform. + wafer_slots: list + Optional list of wafer slots to make solutions for. base_obs_id: str Obs_id to use as a base for matching when merging multiple pointing obs_ids for a wafer. Will default to the pointing obs_id with the greatest number of detectors above the min_R2 threshold. zemax_path: str - If running for a "LAT" tel_type, the path to the zemax file must be specified. + If running for the LAT, the path to the zemax file must be specified. apply_roll: bool Whether or not to apply the obs_id roll angle. Some pointing sets may already be corrected for roll angle. @@ -83,6 +84,8 @@ class SolutionsCfg: (xi_offset, eta_offset) where both are in radians. ufm_to_fp_path: str Path to file that maps wafer_slot to position on focal plane. + fp_to_rx_path: str + Path to file that maps tube slot to position on focal plane. freq_correct_by_muxband: bool If true, apply the same freq offset correction to all resonators in a mux-band. """ @@ -91,7 +94,8 @@ class SolutionsCfg: pointing_results_dir: str results_dir: str wafer_info_path: str - tel_type: str + platform: str + wafer_slots: Optional[list[str]] = None base_obs_id: Optional[str] = None zemax_path: Optional[str] = None apply_roll: bool = True @@ -109,6 +113,8 @@ class SolutionsCfg: initial_pointing_offset: Tuple[float, float] = (0, 0) ufm_to_fp_path: Optional[str] = None + fp_to_rx_path: Optional[str] = None + imprinter_path: Optional[str] = None freq_correct_by_muxband: bool = True ctx: Context = field(init=False) @@ -126,6 +132,9 @@ def from_yaml(cls, path: str) -> "SolutionsCfg": def __post_init__(self): self.ctx = Context(self.ctx_path) + if self.wafer_slots is None: + self.wafer_slots = ["ws.", "ws0", "ws1", "ws2", "ws3", "ws4", "ws5", "ws6"] + if not os.path.exists(self.results_dir): if os.path.exists(os.path.split(self.results_dir)[0]): os.makedirs(self.results_dir) @@ -146,6 +155,14 @@ def __post_init__(self): self.ufm_to_fp_path = os.path.join( self.site_pipeline_cfg_dir, "shared/focalplane/ufm_to_fp.yaml" ) + if self.fp_to_rx_path is None: + self.fp_to_rx_path = os.path.join( + self.site_pipeline_cfg_dir, "shared/focalplane/optics_tubes.yaml" + ) + if self.imprinter_path is None: + self.imprinter_path = os.path.join( + self.site_pipeline_cfg_dir, f"{self.platform}/imprinter.yaml" + ) @dataclass @@ -230,7 +247,7 @@ def pointing_preprocess(cfg: SolutionsCfg, pinfo: PointingInfo): return meta -def merge_pointing_info(cfg: SolutionsCfg, pinfos: List[PointingInfo], base_idx=0): +def merge_pointing_info(cfg: SolutionsCfg, pinfos: List[PointingInfo], base_idx=0, ignore_north_south=False): """ Combine all pointing measurements into a single resonator set, with the median pointing info from all. This requires a base_idx to be specified, @@ -246,7 +263,13 @@ def merge_pointing_info(cfg: SolutionsCfg, pinfos: List[PointingInfo], base_idx= meta = pinfos[base_idx].meta stream_id = meta.det_info.stream_id[0] wafer_slot = meta.det_info.wafer_slot[0] - base_resset = dm.ResSet.from_aman(meta, stream_id=stream_id, pointing=meta.tod_pointing) + + base_resset = dm.ResSet.from_aman( + meta, + stream_id=stream_id, + pointing=meta.tod_pointing, + ignore_north_south=ignore_north_south + ) pointing_map = { r.idx: [(r.xi, r.eta)] for r in base_resset @@ -254,14 +277,19 @@ def merge_pointing_info(cfg: SolutionsCfg, pinfos: List[PointingInfo], base_idx= match_pars = dm.MatchParams( freq_width=cfg.match_pars["pointing"]["freq_width"], - dist_width=np.deg2rad(cfg.match_pars["pointing"]["dist_width"]) + dist_width=np.deg2rad(cfg.match_pars["pointing"]["dist_width"]), ) for i in range(len(pinfos)): if i == base_idx: continue meta = pinfos[i].meta - src = dm.ResSet.from_aman(meta, stream_id=stream_id, pointing=meta.tod_pointing) + src = dm.ResSet.from_aman( + meta, + stream_id=stream_id, + pointing=meta.tod_pointing, + ignore_north_south=ignore_north_south + ) dst = base_resset match = dm.Match(src, dst, match_pars=match_pars) for rsrc, rdst in match.get_match_iter(include_unmatched=False): @@ -390,7 +418,8 @@ def match_wafer( cfg: SolutionsCfg, am: AxisManager, stream_id: str, - meas_rset: Optional[dm.ResSet] + meas_rset: Optional[dm.ResSet], + ignore_north_south: bool = False ) -> MatchSolution: """ Create a match solution for a given wafer slot. @@ -409,19 +438,81 @@ def match_wafer( m = am.det_info.stream_id == stream_id wafer_slot = am.det_info.wafer_slot[m][0] + tube_slot = am.obs_info.tube_slot if meas_rset is None: - src = dm.ResSet.from_aman(am, stream_id, pointing=am[cfg.pointing_field]) + src = dm.ResSet.from_aman( + am, + stream_id, + pointing=am[cfg.pointing_field], + ignore_north_south=ignore_north_south + ) else: src = meas_rset - pt_cfg = dm.PointingConfig( - fp_file=cfg.ufm_to_fp_path, wafer_slot=wafer_slot, tel_type=cfg.tel_type, - zemax_path=cfg.zemax_path, - roll=np.deg2rad(am.obs_info.roll_center) if cfg.apply_roll else 0.0, - tube_slot = am.obs_info.tube_slot - ) - dst = dm.ResSet.from_wafer_info_file(cfg.wafer_info_path, stream_id, pt_cfg=pt_cfg) + if ignore_north_south: + wafer_slots = [] + wafers_dict = {} + + # get wafer and wafer_slot entry from imprinter + with open(cfg.imprinter_path, "r") as f: + imprinter = yaml.safe_load(f) + + for k, v in imprinter['tel_tubes'].items(): + if v['tube_slot'] == tube_slot: + for ws in v['wafer_slots']: + if 'wafer' in ws.keys(): + wafers_dict[ws['wafer_slot']] = (ws['wafer'].split('_')[-1].capitalize()) + wafer_slots.append(ws['wafer_slot']) + else: + wafer_slots = [wafer_slot] + wafers_dict = None + + dst_merged = [] + # loop through and get pointing information for all wafers + for ws in wafer_slots: + pt_cfg = dm.PointingConfig( + fp_file=cfg.ufm_to_fp_path, rx_file=cfg.fp_to_rx_path, + wafer_slot=ws, platform=cfg.platform, + zemax_path=cfg.zemax_path, + roll=np.deg2rad(am.obs_info.roll_center) if cfg.apply_roll else 0.0, + tube_slot = am.obs_info.tube_slot + ) + dst_wafer = dm.ResSet.from_wafer_info_file( + cfg.wafer_info_path, + stream_id, + pt_cfg=pt_cfg, + ignore_north_south=ignore_north_south, + ).as_array() + + # extract current wafer + if wafers_dict is not None: + mask = [wafers_dict[ws] in d["det_id"].astype(str) for d in dst_wafer] + dst_wafer = dst_wafer[mask] + dst_wafer['wafer_slot'][:] = ws + + dst_merged.append(dst_wafer) + dst_merged = np.concatenate(dst_merged) + + # get full wafer info and populate pointing and wafer_slot information + # from merged dst (not necesary if matching a single wafer) + if ignore_north_south: + dst = dm.ResSet.from_wafer_info_file( + cfg.wafer_info_path, + stream_id, + pt_cfg=None, + ignore_north_south=ignore_north_south, + ).as_array() + + merge_fields = ["xi", "eta", "wafer_slot"] + for i, d in enumerate(dst_merged): + idx = np.where(dst['det_id'] == d['det_id'])[0][0] + if idx: + for field in merge_fields: + dst[idx][field] = dst_merged[i][field] + dst = dm.ResSet.from_array(dst) + else: + dst = dm.ResSet.from_array(dst_merged) # first match match_pars = dm.MatchParams( @@ -429,7 +520,7 @@ def match_wafer( dist_width=np.deg2rad(cfg.match_pars["match0"]["dist_width"]), enforce_pointing_reqs=True, allow_unassigned_to_assigned=False, - unassigned_slots=cfg.unassigned_slots + unassigned_slots=cfg.unassigned_slots, ) match = dm.Match(src, dst, match_pars=match_pars, apply_dst_pointing=False) @@ -497,7 +588,6 @@ def match_wafer( match.match_pars.dist_width = np.deg2rad(cfg.match_pars["match2"]["dist_width"]) match._match() - match_iterations.append(deepcopy(match)) return MatchSolution( @@ -555,13 +645,21 @@ def get_wafer_solution( else: base_idx = 0 - meas_rset, pointing_map = merge_pointing_info(cfg, pointing_results, base_idx=base_idx) - # tod_pointing = get_best_tod_pointing(cfg, pointing_results) + # Whether or not to use North/South in match + ignore_north_south = (wafer_slot == "ws.") + + meas_rset, pointing_map = merge_pointing_info( + cfg, pointing_results, base_idx=base_idx, + ignore_north_south=ignore_north_south + ) meta = pointing_results[0].meta stream_id = meta.det_info.stream_id[0] - match_solution = match_wafer(cfg, meta, stream_id, meas_rset=meas_rset) + match_solution = match_wafer( + cfg, meta, stream_id, meas_rset=meas_rset, + ignore_north_south=ignore_north_south + ) solution = FullWaferSolution( pointing_results=pointing_results, @@ -577,8 +675,7 @@ def get_wafer_solution( def solve_all(cfg) -> Dict[str, Optional[FullWaferSolution]]: - wafer_slots = ["ws0", "ws1", "ws2", "ws3", "ws4", "ws5", "ws6"] - results = {ws: get_wafer_solution(cfg, ws, save=True) for ws in wafer_slots} + results = {ws: get_wafer_solution(cfg, ws, save=True) for ws in cfg.wafer_slots} return results diff --git a/sotodlib/coords/optics.py b/sotodlib/coords/optics.py index 01e1ce6c1..dc5c51844 100644 --- a/sotodlib/coords/optics.py +++ b/sotodlib/coords/optics.py @@ -376,7 +376,7 @@ def load_zemax(path): zemax_dat: Dictionairy with data from zemax """ try: - zemax_dat = np.load(path, allow_pickle=True) + zemax_dat = np.load(path, allow_pickle=True, encoding='bytes').items() except Exception as e: logger.error("Can't load data from " + path) raise e @@ -401,6 +401,7 @@ def LAT_optics(zemax_path): zemax_dat = load_zemax(zemax_path) try: LAT = zemax_dat["LAT"][()] + LAT = {k.decode(): v for k, v in zemax_dat["LAT"][()].items()} except Exception as e: logger.error("LAT key missing from dictionary") raise e @@ -439,6 +440,10 @@ def LATR_optics(zemax_path, tube_slot): zemax_dat = load_zemax(zemax_path) try: LATR = zemax_dat["LATR"][()] + LATR = np.array([ + {k.decode(): v for k, v in d.items()} + for d in LATR + ], dtype=object) except Exception as e: logger.error("LATR key missing from dictionary") raise e diff --git a/sotodlib/site_pipeline/update_det_match.py b/sotodlib/site_pipeline/update_det_match.py index 3cf74ca2c..f1560f5dd 100644 --- a/sotodlib/site_pipeline/update_det_match.py +++ b/sotodlib/site_pipeline/update_det_match.py @@ -245,7 +245,9 @@ def get_failed_detsets(cache_file): def run_match_aman(runner: Runner, aman, detset, wafer_slot=None): stream_id = aman.det_info.stream_id[aman.det_info.detset == detset][0] - rs0 = det_match.ResSet.from_aman(aman, stream_id) + ignore_north_south = True if wafer_slot == "ws." else False + + rs0 = det_match.ResSet.from_aman(aman, stream_id, ignore_north_south=ignore_north_south) rs0.name = 'meas' rs1 = load_solution_set(runner, stream_id, wafer_slot=wafer_slot) @@ -261,6 +263,7 @@ def run_match_aman(runner: Runner, aman, detset, wafer_slot=None): match_pars.freq_offset_mhz = opt_freq match = det_match.Match(rs0, rs1, match_pars=match_pars, apply_dst_pointing=runner.cfg.apply_solution_pointing) + return match def run_match(runner: Runner, detset: str) -> bool: @@ -314,23 +317,8 @@ def run_match(runner: Runner, detset: str) -> bool: logger.info(f" - {ds}") for ds in new_detsets: - stream_id = aman.det_info.stream_id[aman.det_info.detset == ds][0] - # Try to get wafer slot info from book idx - if 'wafer_slots' in book_idx: - for ws in book_idx['wafer_slots']: - if ws['stream_id'] == stream_id: - wafer_slot = ws['wafer_slot'] - break - else: - logger.error( - f"Could not find wafer_slot from book index for ds={detset}, " - f"obs_id={obs_id}" - ) - raise Exception("Could not find wafer-slot") - else: - wafer_slot = None - + wafer_slot = aman.det_info.wafer_slot[aman.det_info.detset == ds][0] match = run_match_aman(runner, aman, ds, wafer_slot=wafer_slot) fpath = os.path.join(runner.match_dir, f"{ds}.h5") match.save(fpath)