From 9d5229ff7846b8f7ca5f5967fead9cba351a53f5 Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Wed, 16 Oct 2024 08:23:49 -0400 Subject: [PATCH] Helper function for auto-applying colmap --- src/tdastro/astro_utils/opsim.py | 3 +++ src/tdastro/example_runs/simulate_snia.py | 6 +++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/tdastro/astro_utils/opsim.py b/src/tdastro/astro_utils/opsim.py index f147c3b7..0b2f2986 100644 --- a/src/tdastro/astro_utils/opsim.py +++ b/src/tdastro/astro_utils/opsim.py @@ -145,6 +145,9 @@ def __len__(self): def __getitem__(self, key): """Access the underlying opsim table.""" + # Auto apply the colmap if possible. + if self.colmap is not None and key in self.colmap: + return self.table[self.colmap[key]] return self.table[key] @property diff --git a/src/tdastro/example_runs/simulate_snia.py b/src/tdastro/example_runs/simulate_snia.py index 44b24ab5..6b6db7d4 100644 --- a/src/tdastro/example_runs/simulate_snia.py +++ b/src/tdastro/example_runs/simulate_snia.py @@ -34,11 +34,11 @@ def draw_single_random_sn( obs_index = np.array(opsim.range_search(ra, dec, radius=1.75)) # Update obs_index to only include observations within SN lifespan - phase_obs = opsim[opsim.colmap["time"]][obs_index] - t0 + phase_obs = opsim["time"][obs_index] - t0 obs_index = obs_index[(phase_obs > -20 * (1.0 + z)) & (phase_obs < 50 * (1.0 + z))] - times = opsim[opsim.colmap["time"]][obs_index].to_numpy() - filters = opsim[opsim.colmap["filter"]][obs_index].to_numpy(str) + times = opsim["time"][obs_index].to_numpy() + filters = opsim["filter"][obs_index].to_numpy(str) # Change to match band names in passbands object filters = np.char.add("LSST_", filters)