From 05e9c18785f0550447c2627e97aa7a06d6d4083d Mon Sep 17 00:00:00 2001 From: Bob Ellis Armstrong Date: Tue, 29 Apr 2025 12:44:28 -0700 Subject: [PATCH 1/2] Add code to integrate with SSI in LSST --- notebooks/ssi_with_slslim.ipynb | 288 +++++++++ slsim/LsstSciencePipeline/generate_catalog.py | 588 ++++++++++++++++++ .../LsstSciencePipeline/inject_slsim_base.py | 433 +++++++++++++ .../LsstSciencePipeline/inject_slsim_visit.py | 171 +++++ 4 files changed, 1480 insertions(+) create mode 100644 notebooks/ssi_with_slslim.ipynb create mode 100644 slsim/LsstSciencePipeline/generate_catalog.py create mode 100644 slsim/LsstSciencePipeline/inject_slsim_base.py create mode 100644 slsim/LsstSciencePipeline/inject_slsim_visit.py diff --git a/notebooks/ssi_with_slslim.ipynb b/notebooks/ssi_with_slslim.ipynb new file mode 100644 index 000000000..ff3e0f300 --- /dev/null +++ b/notebooks/ssi_with_slslim.ipynb @@ -0,0 +1,288 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "4e80a982-6590-457a-ab5d-c5727c17847f", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import sys\n", + "import numpy as np\n", + "import subprocess\n", + "from astropy.table import Table\n", + "from astropy.time import Time\n", + "import lsst.daf.butler as dafButler\n", + "from lsst.pipe.base import Pipeline\n", + "%matplotlib ipympl" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aaeba882-202a-48bc-82e6-6b28e348ebfb", + "metadata": {}, + "outputs": [], + "source": [ + "import lsst.afw.display as afwDisplay\n", + "afwDisplay.setDefaultBackend('firefly')\n", + "display1 = afwDisplay.Display(frame=1)\n", + "display2 = afwDisplay.Display(frame=2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b9af51d8-42fd-4125-b210-e335d0e41ec3", + "metadata": {}, + "outputs": [], + "source": [ + "# Catalog generation parameters\n", + "ra_min, ra_max = 52, 54\n", + "dec_min, dec_max = -29, -27\n", + "n_galaxies = 1000\n", + "sky_area_value = 5.0\n", + "\n", + "# Butler configuration\n", + "repo_path = \"/repo/main\" \n", + "instrument = \"LSSTComCam\"\n", + "\n", + "# Specific exposure and detector to process\n", + "visit = 2024110800246\n", + "detector = 5 \n", + "\n", + "# Collection with pvi images\n", + "collection = \"LSSTComCam/runs/DRP/DP1/w_2025_10/DM-49359\"\n", + "\n", + "# Collections for processing results\n", + "inject_collection = \"u/rea3/slsim/injection_test10_calexp\"\n", + "output_collection = \"u/rea3/slsim/test10_calexp\"\n", + "\n", + "# Filter visits to use for ingestion\n", + "filters = [\"r\"]\n", + "\n", + "# Time based selection of visits\n", + "start_time = \"2024-11-01T00:00:00\"\n", + "end_time = \"2024-12-29T00:00:00\"\n", + "\n", + "ngal = 10000\n", + "sky_area = 0.5\n", + "output_file = \"catalog.fits\"\n", + "\n", + "# path to slsim\n", + "slsim_path = \"../slsim/slsim/LsstSciencePipeline/\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c1395cc9-cf4c-4698-a255-3d0498dde63d", + "metadata": {}, + "outputs": [], + "source": [ + "gen_cmd = f\"python {slsim_path}/generate_catalog.py --ra-min {ra_min} --ra-max {ra_max} \"\\\n", + " f\" --dec-min {dec_min} --dec-max {dec_max} --start-time {start_time} --end-time {end_time}\"\\\n", + " f\" --n-galaxies {ngal} --sky-area {sky_area} --repo-path {repo_path}\"\\\n", + " f\" --collection {collection} --instrument {instrument}\"\\\n", + " f\" --filters {\" \".join(filters)} --output-file {output_file}\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "39d287e6-4cb5-4fb8-9509-1d8e157fb7a4", + "metadata": {}, + "outputs": [], + "source": [ + "#os.system(gen_cmd)\n", + "print(gen_cmd)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9ecece51-459b-4726-ae2d-b89f5ba4dd6a", + "metadata": {}, + "outputs": [], + "source": [ + "# Use the ingest_injection_catalog command line tool to ingest the previously created catalog\n", + "injest_cmd = f\"ingest_injection_catalog -b {repo_path} -o {inject_collection}\"\\\n", + " f\" --injection-catalog {output_file} {\" \".join(filters)} -t injection_slsim\"\\\n", + " f\" --format fits\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4645d58b-be42-42aa-ad83-08a6af5eb6a0", + "metadata": {}, + "outputs": [], + "source": [ + "#os.system(injest_cmd)\n", + "print(injest_cmd)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "00db276e-e6a2-4b9b-a56e-c2b108d52c84", + "metadata": {}, + "outputs": [], + "source": [ + "# Create a pipeline configuration file for SLSim injection with DRP processing\n", + "pipeline_yaml = \"config_wfake.yaml\"\n", + "\n", + "with open(pipeline_yaml, \"w\") as f:\n", + " f.write(\"\"\"\n", + "description: Strong lensing source injection pipeline with DRP processing\n", + "instrument: lsst.obs.lsst.LsstComCam\n", + "\n", + "tasks:\n", + " inject_visit_slsim:\n", + " class: slsim.LsstSciencePipeline.inject_slsim_visit.VisitInjectSLSimTask\n", + " config:\n", + " external_psf: false\n", + " external_photo_calib: false\n", + " external_wcs: false\n", + " connections.input_exposure: pvi\n", + " connections.output_exposure: injected_slsim_pvi\n", + " selection: \"np.isin(injection_catalog['visit'], {visit})\"\n", + " subtractImages:\n", + " class: lsst.ip.diffim.subtractImages.AlardLuptonSubtractTask\n", + " config:\n", + " - connections.coaddName: goodSeeing\n", + " connections.template: goodSeeingDiff_templateExp\n", + " connections.science: injected_slsim_pvi\n", + " connections.difference: injected_slsim_goodSeeingDiff_differenceTempExp\n", + " connections.matchedTemplate: injected_slsim_goodSeeingDiff_matchedExp\n", + " detectAndMeasureDiaSources:\n", + " class: lsst.ip.diffim.detectAndMeasure.DetectAndMeasureTask\n", + " config:\n", + " - connections.coaddName: goodSeeing\n", + " connections.diaSources: injected_slsim_goodSeeingDiff_diaSrc\n", + " connections.subtractedMeasuredExposure: injected_slsim_goodSeeingDiff_differenceExp\n", + " connections.science: injected_slsim_pvi\n", + " connections.difference: injected_slsim_goodSeeingDiff_differenceTempExp\n", + " connections.matchedTemplate: injected_slsim_goodSeeingDiff_matchedExp\n", + "\n", + "subsets:\n", + " injected_DRP:\n", + " subset:\n", + " - inject_visit_slsim\n", + " - subtractImages\n", + " - detectAndMeasureDiaSources\n", + " description: >\n", + " Pipeline for injecting strong lensing sources and running difference imaging\n", + "\"\"\")\n", + "\n", + "print(f\"Created pipeline configuration in {pipeline_yaml}\")\n", + "print(\"The pipeline includes strong lensing injection followed by subtraction and detection tasks, exactly matching test.yaml\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26f8ff66-f605-4e15-a8f8-681ff01475b1", + "metadata": {}, + "outputs": [], + "source": [ + "run_cmd = f\"pipetask run -b {repo_path}\"\\\n", + " f\" -d \\\"exposure={visit} and detector={detector}\\\" -p {pipeline_yaml}#injected_DRP\"\\\n", + " f\" -i {collection},{inject_collection}\"\\\n", + " f\" -o {output_collection} --register-dataset-types\"\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "182106f5-287a-4767-9f6b-78aced648ccf", + "metadata": {}, + "outputs": [], + "source": [ + "# os.system(run_cmd)\n", + "print(run_cmd)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "46c94c87-6352-4b98-bfcd-4429af277c35", + "metadata": {}, + "outputs": [], + "source": [ + "butler = dafButler.Butler(repo_path,collections=output_collection)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6ec30488-5a32-4d11-a72f-743d4bb87bed", + "metadata": {}, + "outputs": [], + "source": [ + "injected_pvi = butler.get(\"injected_slsim_pvi\", visit=visit, detector=detector)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2fad76e0-b704-4318-8f46-8e1df3b9ffb7", + "metadata": {}, + "outputs": [], + "source": [ + "diff_pvi = butler.get(\"injected_slsim_goodSeeingDiff_differenceTempExp\", visit=visit, detector=detector)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "91fee920-382b-4951-98dc-3abedad7301b", + "metadata": {}, + "outputs": [], + "source": [ + "display1.mtv(injected_pvi)\n", + "display2.mtv(diff_pvi)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d29c7a4f-e101-42e1-ad57-d5ae07292f97", + "metadata": {}, + "outputs": [], + "source": [ + "output_collection" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fbfb23d6-7c95-45d9-b008-f91723c605b0", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "LSST", + "language": "python", + "name": "lsst" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/slsim/LsstSciencePipeline/generate_catalog.py b/slsim/LsstSciencePipeline/generate_catalog.py new file mode 100644 index 000000000..735f3549e --- /dev/null +++ b/slsim/LsstSciencePipeline/generate_catalog.py @@ -0,0 +1,588 @@ +import argparse +import numpy as np +from astropy.table import Table, vstack +from astropy.time import Time +import lsst.daf.butler as dafButler +from astropy.cosmology import FlatLambdaCDM +from astropy.units import Quantity +from slsim.lens_pop import LensPop +import slsim.Pipelines as pipelines +import slsim.Sources as sources +import slsim.Deflectors as deflectors + +try: + from tqdm import tqdm + + HAS_TQDM = True +except ImportError: + # Create a simple fallback if tqdm is not available + HAS_TQDM = False + + def tqdm(iterable, **kwargs): + """Simple fallback for tqdm progress bar when the library is not available. + + :param iterable: Iterable to iterate over + :type iterable: iterable + :param kwargs: Keyword arguments that would be passed to tqdm (ignored in fallback) + :return: Unchanged input iterable + :rtype: iterable + """ + return iterable + + + + +def generate_master_galaxy_list( + ra_min, ra_max, dec_min, dec_max, n_galaxies=1000, sky_area_value=0.15 +): + """Generate a master list of galaxies (source and lens types) using slsim. + + :param ra_min: Minimum right ascension in degrees + :type ra_min: float + :param ra_max: Maximum right ascension in degrees + :type ra_max: float + :param dec_min: Minimum declination in degrees + :type dec_min: float + :param dec_max: Maximum declination in degrees + :type dec_max: float + :param n_galaxies: Number of galaxies to generate + :type n_galaxies: int + :param sky_area_value: Sky area in square degrees for galaxy simulation + :type sky_area_value: float + :return: Galaxy catalog + :rtype: `astropy.table.Table` + """ + # Define cosmology + cosmo = FlatLambdaCDM(H0=70, Om0=0.3) + + # Define sky area for galaxy simulation and lens population + sky_area = Quantity(value=sky_area_value, unit="deg2") + # For this implementation, we use the same area for both, but they could be different + full_sky_area = Quantity(value=sky_area_value, unit="deg2") + + # Define cuts for galaxy selection + kwargs_deflector_cut = {"band": "i", "band_max": 27, "z_min": 0.1, "z_max": 2} + kwargs_source_cut = {"band": "i", "band_max": 27, "z_min": 0.1, "z_max": 5} + + # Initialize galaxy simulation pipeline + galaxy_simulation_pipeline = pipelines.SkyPyPipeline( + skypy_config=None, sky_area=sky_area, filters=None, cosmo=cosmo + ) + + # Initialize lens population + lens_galaxies = deflectors.EllipticalLensGalaxies( + galaxy_list=galaxy_simulation_pipeline.red_galaxies, + kwargs_cut=kwargs_deflector_cut, + kwargs_mass2light=None, + cosmo=cosmo, + sky_area=sky_area, + ) + + # Initialize source population + source_galaxies = sources.Galaxies( + galaxy_list=galaxy_simulation_pipeline.blue_galaxies, + kwargs_cut=kwargs_source_cut, + cosmo=cosmo, + sky_area=sky_area, + catalog_type="skypy", + ) + + # Create lens population + lenspop = LensPop( + deflector_population=lens_galaxies, + source_population=source_galaxies, + cosmo=cosmo, + sky_area=full_sky_area, + ) + + # Get lens and source galaxies + lens = lenspop._lens_galaxies._galaxy_select + source = lenspop._sources._galaxy_select + + # Check if we have enough galaxies for the requested catalog size + if len(lens) < n_galaxies or len(source) < n_galaxies: + raise ValueError( + f"Not enough galaxies generated. Requested {n_galaxies} of each type, but " + f"only got {len(lens)} lens galaxies and {len(source)} source galaxies. " + f"Try increasing sky_area_value or decreasing n_galaxies." + ) + + # Process lens galaxies + lens_ell_mask = lens["n_sersic"] < -0.999 + n_ell_lens = np.sum(lens_ell_mask) + + if n_ell_lens > 0: + phi = np.random.uniform(0, np.pi, size=n_ell_lens) + e = lens["ellipticity"][lens_ell_mask].data + ep = (1 - np.sqrt(1 - e**2)) / e + e1 = ep * np.cos(2 * phi) + e2 = ep * np.sin(2 * phi) + lens["e1_light"][lens_ell_mask] = e1 + lens["e2_light"][lens_ell_mask] = e2 + lens["n_sersic"][lens_ell_mask] = 4 + + # Process source galaxies + source_ell_mask = source["n_sersic"] < -0.999 + n_ell_source = np.sum(source_ell_mask) + + if n_ell_source > 0: + phi = np.random.uniform(0, np.pi, size=n_ell_source) + e = source["ellipticity"][source_ell_mask].data + ep = (1 - np.sqrt(1 - e**2)) / e + e1 = ep * np.cos(2 * phi) + e2 = ep * np.sin(2 * phi) + source["e1"][source_ell_mask] = e1 + source["e2"][source_ell_mask] = e2 + source["n_sersic"][source_ell_mask] = 1 + + + # Generate random positions within the specified RA/Dec box + ra = np.random.uniform(ra_min, ra_max, size=n_galaxies) + dec = np.random.uniform(dec_min, dec_max, size=n_galaxies) + + # Initialize ID, RA, Dec fields + source["ra"] = np.nan + source["dec"] = np.nan + source["id"] = -1 + source["type"] = "source" + source["source_type"] = "extended" + source["light_profile"] = "single_sersic" + lens["ra"] = np.nan + lens["dec"] = np.nan + lens["id"] = -1 + lens["type"] = "lens" + lens["deflector_type"] = "EPL" + + # Randomly select galaxies and assign positions + lens_index = np.random.choice(range(len(lens)), n_galaxies, replace=False) + source_index = np.random.choice(range(len(source)), n_galaxies, replace=False) + + source["id"][source_index] = np.arange(0, n_galaxies).astype(int) + lens["id"][lens_index] = np.arange(0, n_galaxies).astype(int) + source["ra"][source_index] = ra + source["dec"][source_index] = dec + lens["ra"][lens_index] = ra + lens["dec"][lens_index] = dec + + # Combine sources and lenses and expand coefficient array + gals = vstack([source[source_index], lens[lens_index]]) + + # Handle coefficient array expansion + if "coeff" in gals.colnames: + if len(gals["coeff"].shape) > 1 and gals["coeff"].shape[1] >= 5: + gals["coeff0"] = gals["coeff"][:, 0] + gals["coeff1"] = gals["coeff"][:, 1] + gals["coeff2"] = gals["coeff"][:, 2] + gals["coeff3"] = gals["coeff"][:, 3] + gals["coeff4"] = gals["coeff"][:, 4] + del gals["coeff"] + else: + # Handle the case when coeffs don't have expected shape + for i in range(5): + if f"coeff{i}" not in gals.colnames: + gals[f"coeff{i}"] = np.zeros(len(gals)) + else: + # If coeffs don't exist, create empty ones + for i in range(5): + gals[f"coeff{i}"] = np.zeros(len(gals)) + + return gals + + +def get_calexps_in_region( + butler, + collection, + time_range, + ra_range, + dec_range, + instrument="LSSTComCam", + filters=["u", "g", "r", "i", "z", "y"], + max_calexps=10, +): + """Get calexp objects from the butler within given time and spatial constraints. + + :param butler: The butler instance + :type butler: `lsst.daf.butler.Butler` + :param collection: The collection to query + :type collection: str + :param time_range: (start_time, end_time) as astropy.time.Time objects + :type time_range: tuple + :param ra_range: (ra_min, ra_max) in degrees + :type ra_range: tuple + :param dec_range: (dec_min, dec_max) in degrees + :type dec_range: tuple + :param instrument: The instrument name to query + :type instrument: str + :param filters: List of filters to include + :type filters: list + :param max_calexps: Maximum number of calexps to process + :type max_calexps: int + :return: Tuple containing visit_list, visit_ras, visit_decs, visit_times + :rtype: tuple + """ + # Query for available visit IDs with the filter constraint + filter_conditions = [] + for f in filters: + filter_conditions.append(f"band='{f}'") + + filter_clause = " OR ".join(filter_conditions) + where_clause = f"instrument='{instrument}' AND ({filter_clause})" + + print(f"Querying visits with where clause: {where_clause}") + try: + # First, query for visits + visitIds = butler.registry.queryDataIds( + ["visit"], + datasets="visitSummary", + collections=collection, + where=where_clause, + ) + print(f"Successfully queried visit IDs") + except Exception as e: + print(f"Error querying visits: {e}") + return [], [], [], [] + + # Get visitSummaries and filter by time and spatial overlap + time_start = time_range[0] + time_end = time_range[1] + + # Convert visitIds to a list for progress tracking + visit_list = list(visitIds) + total_visits = len(visit_list) + print(f"Found {total_visits} visits to check") + + # Process each visit summary + visit_ras = [] + visit_decs = [] + visit_times = [] + + for visit_dataId in tqdm(visit_list, desc="Checking visitSummaries", unit="visit"): + try: + # Get the visitSummary + visit_id = visit_dataId["visit"] + visitSummary = butler.get( + "visitSummary", visit_dataId, collections=collection + ) + + # Get the date from day_obs + try: + day_obs = visit_dataId["day_obs"] + visit_str = str(visit_dataId["visit"]) + print(visit_str) + if len(visit_str) >= 14: # Full timestamp format + time_part = visit_str[-6:] # Extract HHMMSS + hour = int(time_part[:2]) + minute = int(time_part[2:4]) + second = int(time_part[4:6]) + # Create date string in ISO format + date_str = f"{day_obs // 10000:04d}-{(day_obs // 100) % 100:02d}-{day_obs % 100:02d}T{hour:02d}:{minute:02d}:{second:02d}" + visit_date = Time(date_str, format="isot") + else: + # Just use the day with time 00:00:00 + date_str = f"{day_obs // 10000:04d}-{(day_obs // 100) % 100:02d}-{day_obs % 100:02d}T00:00:00" + visit_date = Time(date_str, format="isot") + + # Check if it's in our time range + if visit_date < time_start or visit_date > time_end: + continue + except Exception: + continue + + # Filter already checked in the query + + # We'll check each detector individually for overlap with our region + + # Check each detector in the visit for overlap with our region + detector_ids = visitSummary["id"] + + # Initialize min/max values for this visit + min_ra_visit = np.inf + max_ra_visit = -np.inf + min_dec_visit = np.inf + max_dec_visit = -np.inf + + try: + for i, detector_id in enumerate(detector_ids): + ra_corners = visitSummary["raCorners"][i] + dec_corners = visitSummary["decCorners"][i] + min_ra, max_ra = min(ra_corners), max(ra_corners) + min_dec, max_dec = min(dec_corners), max(dec_corners) + + # Update visit-level min/max values + min_ra_visit = min(min_ra, min_ra_visit) + max_ra_visit = max(max_ra, max_ra_visit) + min_dec_visit = min(min_dec, min_dec_visit) + max_dec_visit = max(max_dec, max_dec_visit) + except Exception: + continue + + # Add this visit's data to our lists + visit_ras.extend((min_ra_visit, max_ra_visit)) + visit_decs.extend((min_dec_visit, max_dec_visit)) + visit_times.append(visit_date) + + except Exception: + pass + + return visit_list, visit_ras, visit_decs, visit_times + + +def find_galaxies_for_each_visit( + visit_list, visit_ras, visit_decs, visit_times, galaxy_catalog, reference_time=None +): + """Find which galaxies from the master catalog overlap with each visit's boundaries and record the visit ID and time information. + + :param visit_list: List of visit dataIds + :type visit_list: list + :param visit_ras: List of RA values (min/max pairs) for visits + :type visit_ras: list + :param visit_decs: List of Dec values (min/max pairs) for visits + :type visit_decs: list + :param visit_times: List of visit timestamps + :type visit_times: list + :param galaxy_catalog: The master galaxy catalog + :type galaxy_catalog: `astropy.table.Table` + :param reference_time: Reference time to compute relative times + :type reference_time: `astropy.time.Time` + :return: Table containing galaxies with their overlapping visits + :rtype: `astropy.table.Table` + """ + # Prepare to store all galaxies with their visit information + all_overlaps = [] + + # Set reference time if not provided + if reference_time is None and visit_times: + reference_time = min(visit_times) + + # Process each visit + for i, (visit_dataId, visit_time) in enumerate(zip(visit_list, visit_times)): + visit_id = visit_dataId["visit"] + + # Calculate index into the ra/dec lists, which have min/max pairs for each visit + # Each visit has 2 entries (min/max) in the ra/dec lists + idx = i * 2 + + # Get the RA/Dec bounds for this visit + # Make sure we don't go beyond list bounds + if idx + 1 < len(visit_ras) and idx + 1 < len(visit_decs): + min_ra = visit_ras[idx] + max_ra = visit_ras[idx + 1] + min_dec = visit_decs[idx] + max_dec = visit_decs[idx + 1] + else: + continue # Skip if indices are out of bounds + + # Find galaxies within these bounds + mask = ( + (galaxy_catalog["ra"] >= min_ra) + & (galaxy_catalog["ra"] <= max_ra) + & (galaxy_catalog["dec"] >= min_dec) + & (galaxy_catalog["dec"] <= max_dec) + ) + + visit_galaxies = galaxy_catalog[mask].copy() + + # Skip if no galaxies found + if len(visit_galaxies) == 0: + continue + + # Add visit information to each galaxy + visit_galaxies["visit"] = visit_id + + # Store observation time as ISO format string + visit_galaxies["observation_time"] = visit_time.isot + + # Add relative time if reference_time is provided + if reference_time is not None: + time_delta = (visit_time - reference_time).sec + visit_galaxies["time_delta"] = time_delta + + all_overlaps.append(visit_galaxies) + + # Combine all visit-galaxy pairs + if all_overlaps: + return vstack(all_overlaps) + else: + # Return empty table with same structure + if len(galaxy_catalog) > 0: + empty_table = galaxy_catalog[:0].copy() + empty_table["visit"] = [] + # Create observation_time as string column + empty_table["observation_time"] = np.array( + [], dtype="S26" + ) # ISO format datetime string + if reference_time is not None: + empty_table["time_delta"] = [] + return empty_table + else: + return Table() # Empty table + + +def parse_args(): + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + description="Generate a catalog of galaxies and find overlaps with LSST visits." + ) + + # Region of interest + parser.add_argument("--ra-min", type=float, default=52, help="Minimum right ascension in degrees") + parser.add_argument("--ra-max", type=float, default=54, help="Maximum right ascension in degrees") + parser.add_argument("--dec-min", type=float, default=-29, help="Minimum declination in degrees") + parser.add_argument("--dec-max", type=float, default=-27, help="Maximum declination in degrees") + + # Time range + parser.add_argument( + "--start-time", + type=str, + default="2024-11-01T00:00:00", + help="Start time in ISO format (YYYY-MM-DDTHH:MM:SS)" + ) + parser.add_argument( + "--end-time", + type=str, + default="2024-12-29T00:00:00", + help="End time in ISO format (YYYY-MM-DDTHH:MM:SS)" + ) + + # Galaxy generation parameters + parser.add_argument( + "--n-galaxies", + type=int, + default=10000, + help="Number of galaxies to generate" + ) + parser.add_argument( + "--sky-area", + type=float, + default=5.0, + help="Sky area in square degrees for galaxy simulation" + ) + + # Butler parameters + parser.add_argument( + "--repo-path", + type=str, + default="/repo/main", + help="Path to the butler repository" + ) + parser.add_argument( + "--collection", + type=str, + default="LSSTComCam/runs/DRP/DP1/w_2025_07/DM-48940", + help="Butler collection to query" + ) + parser.add_argument( + "--instrument", + type=str, + default="LSSTComCam", + help="Instrument name to query" + ) + parser.add_argument( + "--filters", + type=str, + nargs="+", + default=["r"], + help="List of filters to include" + ) + + # Output + parser.add_argument( + "--output-file", + type=str, + default="galaxy_visit_overlaps.fits", + help="Output file name for galaxy-visit overlaps" + ) + + return parser.parse_args() + + +def main(): + """Main function to process calexps and find overlapping galaxies.""" + args = parse_args() + + # Define region of interest + ra_min, ra_max = args.ra_min, args.ra_max + dec_min, dec_max = args.dec_min, args.dec_max + + # Define time range + start_time = Time(args.start_time, format="isot") + end_time = Time(args.end_time, format="isot") + + # Generate master galaxy catalog with sky area parameter + print("Generating master galaxy catalog...") + sky_area_value = args.sky_area + n_galaxies = args.n_galaxies + + master_galaxies = generate_master_galaxy_list( + ra_min=ra_min, + ra_max=ra_max, + dec_min=dec_min, + dec_max=dec_max, + n_galaxies=n_galaxies, + sky_area_value=sky_area_value, + ) + print(f"Generated {len(master_galaxies)} galaxy entries in catalog") + + # Connect to the butler repository + print("Connecting to butler repository...") + repo_path = args.repo_path + butler = dafButler.Butler(repo_path) + + # Get visit information + print("Fetching visit information in the specified region and time range...") + collection = args.collection + instrument = args.instrument + filters = args.filters + print(f"Limiting search to filters: {filters}") + + visit_list, visit_ras, visit_decs, visit_times = get_calexps_in_region( + butler, + collection, + (start_time, end_time), + (ra_min, ra_max), + (dec_min, dec_max), + instrument=instrument, + filters=filters, + ) + + print(f"Found {len(visit_list)} visits with valid time and spatial information") + + if not visit_times: + print("No visits found. Exiting.") + return + + # Find reference time + reference_time = min(visit_times) if visit_times else None + if reference_time: + print(f"Reference time set to: {reference_time}") + else: + print("No reference time available") + return + + # Find galaxies overlapping with each visit and add visit information + overlapping_galaxies = find_galaxies_for_each_visit( + visit_list, visit_ras, visit_decs, visit_times, master_galaxies, reference_time + ) + + # Count unique galaxies and total overlaps + unique_galaxy_count = ( + len(set(overlapping_galaxies["id"])) if len(overlapping_galaxies) > 0 else 0 + ) + total_overlaps = len(overlapping_galaxies) + + print( + f"Found {unique_galaxy_count} unique galaxies with {total_overlaps} total overlaps across visits" + ) + + # Save results + if total_overlaps > 0: + # Save the results + output_file = args.output_file + overlapping_galaxies.write(output_file, overwrite=True) + print(f"Saved {total_overlaps} galaxy-visit overlaps to {output_file}") + else: + print("No overlapping galaxies found") + + +if __name__ == "__main__": + main() diff --git a/slsim/LsstSciencePipeline/inject_slsim_base.py b/slsim/LsstSciencePipeline/inject_slsim_base.py new file mode 100644 index 000000000..74ffa8923 --- /dev/null +++ b/slsim/LsstSciencePipeline/inject_slsim_base.py @@ -0,0 +1,433 @@ +from __future__ import annotations + +__all__ = ["BaseInjectSLSimConnections", "BaseInjectSLSimConfig", "BaseInjectSLSimTask"] + +from typing import cast + +import numpy as np + +import lsst.geom as geom +from astropy.table import Table, hstack, vstack +from lsst.pex.config import Field +from lsst.pipe.base.connectionTypes import PrerequisiteInput +from lsst.pex.exceptions import InvalidParameterError +from lsst.pipe.base import Struct + +from astropy.cosmology import FlatLambdaCDM + +from lsst.source.injection.inject_base import ( + BaseInjectConnections, + BaseInjectConfig, + BaseInjectTask, +) + +from slsim.Sources.source import Source +from slsim.Deflectors.deflector import Deflector +from slsim.LOS.los_pop import LOSPop +from slsim.lens import Lens +from slsim.image_simulation import lens_image + + +class BaseInjectSLSimConnections( + BaseInjectConnections, + dimensions=("instrument",), + defaultTemplates={ + "injection_prefix": "injection_slsim_", + "injected_prefix": "injected_slsim_", + }, +): + """Base connections for strong lensing source injection tasks. + + This class extends the BaseInjectConnections class from LSST's source injection + framework. For general information on how to use the injection framework, refer + to the LSST source injection classes documentation at: + https://github.com/lsst/source-injection + """ + + injection_catalogs = PrerequisiteInput( + doc="Set of catalogs of sources to draw inputs from.", + # name="{injection_prefix}catalog", + name="injection_slsim", + dimensions=("htm7", "band"), + storageClass="ArrowAstropy", + minimum=0, + multiple=True, + ) + + +class BaseInjectSLSimConfig( + BaseInjectConfig, pipelineConnections=BaseInjectSLSimConnections +): + """Base configuration for strong lensing source injection tasks. + + This class extends the BaseInjectConfig class from LSST's source injection + framework and customizes it for strong lensing injection. For general + information on how to use the injection framework's configuration options, + refer to the LSST source injection documentation at: + https://github.com/lsst/source-injection + """ + + # Catalog manipulation options. + process_all_data_ids = Field[bool]( + doc="If True, all input data IDs will be processed, even those where no synthetic sources were " + "identified for injection. In such an eventuality this returns a clone of the input image, renamed " + "to the *output_exposure* connection name and with an empty *mask_plane_name* mask plane attached.", + default=False, + ) + trim_padding = Field[int]( + doc="Size of the pixel padding surrounding the image. Only those synthetic sources with a centroid " + "falling within the ``image + trim_padding`` region will be considered for source injection.", + default=100, + optional=True, + ) + selection = Field[str]( + doc="A string that can be evaluated as a boolean expression to select rows in the input injection " + "catalog. To make use of this configuration option, the internal object name ``injection_catalog`` " + "must be used. For example, to select all sources with a magnitude in the range 20.0 < mag < 25.0, " + "set ``selection=\"(injection_catalog['mag'] > 20.0) & (injection_catalog['mag'] < 25.0)\"``. " + "The ``{visit}`` field will be substituted for the current visit ID of the exposure being processed. " + "For example, to select only visits that match a user-supplied visit column in the input injection " + "catalog, set ``selection=\"np.isin(injection_catalog['visit'], {visit})\"``.", + optional=True, + ) + # General configuration options. + mask_plane_name = Field[str]( + doc="Name assigned to the injected mask plane which is attached to the output exposure.", + default="SLSIM_INJECTED", + ) + # Size of injected pixels, same for all + stamp_size = Field[int]( + doc="Size of stamp for each injected source", + default=60, + optional=True, + ) + + def setDefaults(self): + """Set defaults for configuration parameters. + + :return: None + """ + super().setDefaults() + + +class BaseInjectSLSimTask(BaseInjectTask): + """Base class for injecting strong lensing sources into images. + + This class extends the BaseInjectTask class from LSST's source injection + framework to handle the specific requirements of strong lensing image simulation + using the SLSim package. It uses the same overall workflow as the base class, + but adds specialized handling for lens-source galaxy pairs and their images. + + For information on the general source injection framework and how to use it, + refer to the LSST source injection documentation at: + https://github.com/lsst/source-injection + + For details on the SLSim strong lensing simulation package: + https://github.com/LSST-strong-lensing/slsim + """ + + _DefaultName = "baseInjectSLSimTask" + ConfigClass = BaseInjectSLSimConfig + + def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs): + """Inject strong lensing sources into an image. + + :param injection_catalogs: Tract level injection catalogs that potentially cover the named input exposure + :type injection_catalogs: `list` [`astropy.table.Table`] + :param input_exposure: The exposure sources will be injected into + :type input_exposure: `lsst.afw.image.ExposureF` + :param psf: PSF model + :type psf: `lsst.meas.algorithms.ImagePsf` + :param photo_calib: Photometric calibration used to calibrate injected sources + :type photo_calib: `lsst.afw.image.PhotoCalib` + :param wcs: WCS used to calibrate injected sources + :type wcs: `lsst.afw.geom.SkyWcs` + :return: Struct containing output_exposure and output_catalog + :rtype: `lsst.pipe.base.Struct` with output_exposure (`lsst.afw.image.ExposureF`) and output_catalog (`lsst.afw.table.SourceCatalog`) + """ + self.config = cast(BaseInjectSLSimConfig, self.config) + + # Make empty table if none supplied to support process_all_data_ids. + if len(injection_catalogs) == 0: + if self.config.process_all_data_ids: + injection_catalogs = [Table(names=["ra", "dec"])] + else: + raise RuntimeError( + "No injection sources overlap the data query. Check injection catalog coverage." + ) + + # Consolidate injection catalogs and compose main injection catalog. + injection_catalog = self._compose_injection_catalog(injection_catalogs) + + # Clean the injection catalog of sources which are not injectable. + injection_catalog = self._clean_sources(injection_catalog, input_exposure) + + # Injection binary flag lookup dictionary. + binary_flags = { + "SLSIM_FAILURE": 0, + "NOT_FULL_OVERLAP": 1, + "PSF_COMPUTE_ERROR": 2, + } + + # Check that sources in the injection catalog are able to be injected. + injection_catalog = self._check_sources(injection_catalog, binary_flags) + + # Inject sources into input_exposure. + good_injections: list[bool] = injection_catalog["injection_flag"] == 0 + good_injections_index = [i for i, val in enumerate(good_injections) if val] + num_injection_sources = np.sum(good_injections) + calib_flux_radius = None + num_pix = self.config.stamp_size + + if num_injection_sources > 0: + cosmo = FlatLambdaCDM(H0=70, Om0=0.3) + + exp_bbox = input_exposure.getBBox() + zeropoint = 2.5 * np.log10(photo_calib.getInstFluxAtZeroMagnitude()) + band = input_exposure.getInfo().getFilter().bandLabel + exp_time = input_exposure.getInfo().getVisitInfo().getExposureTime() + + lens_mask = injection_catalog[good_injections]["type"] == "lens" + source_mask = injection_catalog[good_injections]["type"] == "source" + lens_index = injection_catalog[good_injections][lens_mask]["id"].data + for li in lens_index: + lens_sel = injection_catalog[good_injections][lens_mask]["id"] == li + source_sel = injection_catalog[good_injections][source_mask]["id"] == li + lens_dict = Table( + injection_catalog[good_injections][lens_mask][lens_sel] + ) + source_dict = Table( + injection_catalog[good_injections][source_mask][source_sel] + ) + + sky = geom.SpherePoint( + lens_dict["ra"][0], lens_dict["dec"][0], geom.degrees + ) + xy = wcs.skyToPixel(sky) + + try: + psf_xy = psf.computeKernelImage(xy).array + except InvalidParameterError: + injection_catalog[good_injections][lens_mask]["injection_flag"] += ( + 2 ** binary_flags["PSF_COMPUTE_ERROR"] + ) + injection_catalog[good_injections][source_mask][ + "injection_flag" + ] += (2 ** binary_flags["PSF_COMPUTE_ERROR"]) + continue + + pixscale = wcs.getPixelScale(xy).asArcseconds() + bbox = geom.Box2I( + geom.Point2I(xy.getX() - num_pix / 2, xy.getY() - num_pix / 2), + geom.Extent2I(num_pix, num_pix), + ) + + if exp_bbox.contains(bbox) == False: + injection_catalog[good_injections][lens_mask]["injection_flag"] += ( + 2 ** binary_flags["NOT_FULL_OVERLAP"] + ) + injection_catalog[good_injections][source_mask][ + "injection_flag" + ] += (2 ** binary_flags["NOT_FULL_OVERLAP"]) + continue + + if self.config.calib_flux_radius is not None: + apCorr = psf.computeApertureFlux(self.config.calib_flux_radius, xy) + psf_ker = psf_xy / apCorr + else: + psf_ker = psf_xy + + matrix = ( + wcs.linearizePixelToSky(sky, geom.arcseconds) + .getLinear() + .getMatrix() + ) + + # SLSim has problems with matrix that include reflections like + # because it assumes det is positive, need to recreate matrix + scale = np.sqrt(abs(np.linalg.det(matrix))) + matrix = np.array([[scale, 0], [0, scale]]) + + # Consolidate vector objects because butler stored them separately + source_dict["coeff"] = np.array( + [ + [ + source_dict["coeff0"], + source_dict["coeff1"], + source_dict["coeff2"], + source_dict["coeff3"], + source_dict["coeff4"], + ] + ] + ) + lens_dict["coeff"] = np.array( + [ + [ + lens_dict["coeff0"], + lens_dict["coeff1"], + lens_dict["coeff2"], + lens_dict["coeff3"], + lens_dict["coeff4"], + ] + ] + ) + + try: + source = Source( + source_dict=source_dict, + cosmo=cosmo, + source_type=source_dict["source_type"][0], + light_profile=source_dict["light_profile"][0], + ) + lens = Deflector( + deflector_type=lens_dict["deflector_type"][0], + deflector_dict=lens_dict, + ) + los_pop = LOSPop() + los = los_pop.draw_los( + source_redshift=source.redshift, + deflector_redshift=lens.redshift, + ) + + lens_class = Lens( + source_class=source, + deflector_class=lens, + cosmo=cosmo, + los_class=los, + ) + + lens_im = lens_image( + lens_class=lens_class, + band=band, + mag_zero_point=zeropoint, + num_pix=num_pix, + psf_kernel=psf_ker, + transform_pix2angle=matrix, + exposure_time=exp_time, + ) + except InvalidParameterError: + injection_catalog[good_injections][lens_mask]["injection_flag"] += ( + 2 ** binary_flags["SLSIM_FAILURE"] + ) + injection_catalog[good_injections][source_mask][ + "injection_flag" + ] += (2 ** binary_flags["SLSIM_FAILURE"]) + except Exception as e: + print("Failed: ", e) + continue + + if np.sum(lens_im) > 0: + input_exposure.mask.addMaskPlane(self.config.mask_plane_name) + bitvalue = input_exposure.mask.getPlaneBitMask( + self.config.mask_plane_name + ) + input_exposure[bbox].mask.array |= bitvalue + + input_exposure[bbox].image.array += lens_im + + # Add injection provenance and injection flags metadata. + metadata = input_exposure.getMetadata() + input_dataset_type = self.config.connections.input_exposure.format( + **self.config.connections.toDict() + ) + metadata.set( + "SLSIM_INJECTED", + input_dataset_type, + "Initial source injection dataset type", + ) + for flag, value in sorted(binary_flags.items(), key=lambda item: item[1]): + injection_catalog.meta[flag] = value + + output_struct = Struct( + output_exposure=input_exposure, output_catalog=injection_catalog + ) + return output_struct + + def _compose_injection_catalog(self, injection_catalogs): + """Consolidate injection catalogs and compose main injection catalog. + + If multiple injection catalogs are input, all catalogs are + concatenated together. + + A running injection_id, specific to this dataset ref, is assigned to + each source in the output injection catalog if not provided. + + :param injection_catalogs: Set of synthetic source catalogs to concatenate + :type injection_catalogs: `list` [`astropy.table.Table`] + :return: Catalog of sources to be injected + :rtype: `astropy.table.Table` + """ + self.config = cast(BaseInjectConfig, self.config) + + # Generate injection IDs (if not provided) and injection flag column. + injection_data = vstack(injection_catalogs) + if "injection_id" in injection_data.columns: + injection_id = injection_data["injection_id"] + injection_data.remove_column("injection_id") + else: + injection_id = range(len(injection_data)) + injection_header = Table( + { + "injection_id": injection_id, + "injection_flag": np.zeros(len(injection_data), dtype=int), + } + ) + + # Construct final injection catalog. + injection_catalog = hstack([injection_header, injection_data]) + + # Log and return. + num_injection_catalogs = np.sum( + [len(table) > 0 for table in injection_catalogs] + ) + grammar1 = "source" if len(injection_catalog) == 1 else "sources" + grammar2 = "trixel" if num_injection_catalogs == 1 else "trixels" + self.log.info( + "Retrieved %d injection %s from %d HTM %s.", + len(injection_catalog), + grammar1, + num_injection_catalogs, + grammar2, + ) + return injection_catalog + + def _check_sources(self, injection_catalog, binary_flags): + """Check that sources in the injection catalog are able to be injected. + + This method will check that sources in the injection catalog are able + to be injected, and will flag them if not. Checks will be made on a + number of parameters, including magnitude, source type and Sérsic index + (where relevant). + + Legacy profile types will be renamed to their standardized GalSim + equivalents; any source profile types that are not GalSim classes will + be flagged. + + Note: Unlike the cleaning method, no sources are actually removed here. + Instead, a binary flag is set in the *injection_flag* column for each + source. Only unflagged sources will be generated for source injection. + + :param injection_catalog: Catalog of sources to be injected + :type injection_catalog: `astropy.table.Table` + :param binary_flags: Dictionary of binary flags to be used in the injection_flag column + :type binary_flags: `dict` [`str`, `int`] + :return: The cleaned catalog of sources to be injected + :rtype: `astropy.table.Table` + """ + self.config = cast(BaseInjectConfig, self.config) + + # Exit early if there are no sources to inject. + if len(injection_catalog) == 0: + self.log.info("Catalog checking not applied to empty injection catalog.") + return injection_catalog + + num_flagged_total = np.sum(injection_catalog["injection_flag"] != 0) + grammar = "source" if len(injection_catalog) == 1 else "sources" + self.log.info( + "Catalog checking flagged %d of %d %s; %d remaining for source generation.", + num_flagged_total, + len(injection_catalog), + grammar, + np.sum(injection_catalog["injection_flag"] == 0), + ) + return injection_catalog diff --git a/slsim/LsstSciencePipeline/inject_slsim_visit.py b/slsim/LsstSciencePipeline/inject_slsim_visit.py new file mode 100644 index 000000000..5b9884d47 --- /dev/null +++ b/slsim/LsstSciencePipeline/inject_slsim_visit.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +__all__ = [ + "VisitInjectSLSimConnections", + "VisitInjectSLSimConfig", + "VisitInjectSLSimTask", +] + +from typing import cast + +from lsst.pex.config import Field +from lsst.pipe.base.connectionTypes import Input, Output + +from .inject_slsim_base import ( + BaseInjectSLSimConfig, + BaseInjectSLSimConnections, + BaseInjectSLSimTask, +) + + +class VisitInjectSLSimConnections( # type: ignore [call-arg] + BaseInjectSLSimConnections, + dimensions=("instrument", "visit", "detector"), +): + """Visit-level connections for strong lensing source injection tasks. + + This class extends BaseInjectSLSimConnections to handle visit-level injections. + It configures the appropriate input and output connections for the visit-level + injection tasks. + + For information on the LSST source injection framework this is built upon: + https://github.com/lsst/source-injection + """ + + visit_summary = Input( + doc="A visit summary table containing PSF, PhotoCalib and WCS information.", + name="finalVisitSummary", + storageClass="ExposureCatalog", + dimensions=("visit",), + deferLoad=True, + ) + input_exposure = Input( + doc="Exposure to inject synthetic sources into.", + name="calexp", + storageClass="ExposureF", + dimensions=("instrument", "visit", "detector"), + ) + output_exposure = Output( + doc="Injected Exposure.", + name="{injected_prefix}calexp", + storageClass="ExposureF", + dimensions=("instrument", "visit", "detector"), + ) + output_catalog = Output( + doc="Catalog of injected sources.", + name="{injected_prefix}calexp_catalog", + storageClass="ArrowAstropy", + dimensions=("instrument", "visit", "detector"), + ) + + def __init__(self, *, config=None): + """Initialize visit-level connection parameters. + + :param config: Configuration for the connection + :type config: `VisitInjectSLSimConfig` + :return: None + """ + config = cast(VisitInjectSLSimConfig, config) + + super().__init__(config=config) + if ( + not config.external_psf + and not config.external_photo_calib + and not config.external_wcs + ): + self.inputs.remove("visit_summary") + + +class VisitInjectSLSimConfig( # type: ignore [call-arg] + BaseInjectSLSimConfig, + pipelineConnections=VisitInjectSLSimConnections, +): + """Visit-level configuration for strong lensing source injection tasks. + + This class extends BaseInjectSLSimConfig to provide visit-level configuration + options for strong lensing source injection. For information on the LSST + source injection framework this is built upon: + https://github.com/lsst/source-injection + """ + + # Calibrated data options. + external_psf = Field[bool]( + doc="If True, use the PSF model from a visit summary table. " + "If False (default), use the PSF model attached to the input exposure.", + dtype=bool, + default=False, + ) + external_photo_calib = Field[bool]( + doc="If True, use the photometric calibration from a visit summary table. " + "If False (default), use the photometric calibration attached to the input exposure.", + dtype=bool, + default=False, + ) + external_wcs = Field[bool]( + doc="If True, use the astrometric calibration from a visit summary table. " + "If False (default), use the astrometric calibration attached to the input exposure.", + dtype=bool, + default=False, + ) + + +class VisitInjectSLSimTask(BaseInjectSLSimTask): + """Visit-level class for injecting strong lensing sources into images. + + This task extends BaseInjectSLSimTask to operate at the visit level, handling + visit-specific data structures and metadata. It provides functionality to + inject strong lensing sources into visit-level exposures. + + For information on the LSST source injection framework this is built upon: + https://github.com/lsst/source-injection + """ + + _DefaultName = "visitInjectSLSimTask" + ConfigClass = VisitInjectSLSimConfig + + def runQuantum(self, butler_quantum_context, input_refs, output_refs): + """Run the task on a quantum of data. + + :param butler_quantum_context: Butler quantum context + :type butler_quantum_context: `lsst.daf.butler.QuantumContext` + :param input_refs: Input dataset references + :type input_refs: `dict` + :param output_refs: Output dataset references + :type output_refs: `dict` + :return: None + """ + inputs = butler_quantum_context.get(input_refs) + detector_id = inputs["input_exposure"].getDetector().getId() + + try: + visit_summary = inputs["visit_summary"].get() + except KeyError: + # Use internal PSF, PhotoCalib and WCS. + inputs["psf"] = inputs["input_exposure"].getPsf() + inputs["photo_calib"] = inputs["input_exposure"].getPhotoCalib() + inputs["wcs"] = inputs["input_exposure"].getWcs() + else: + # Use external PSF, PhotoCalib and WCS. + detector_summary = visit_summary.find(detector_id) + if detector_summary: + inputs["psf"] = detector_summary.getPsf() + inputs["photo_calib"] = detector_summary.getPhotoCalib() + inputs["wcs"] = detector_summary.getWcs() + else: + raise RuntimeError( + f"No record for detector {detector_id} found in visit summary table." + ) + + input_keys = [ + "injection_catalogs", + "input_exposure", + "sky_map", + "psf", + "photo_calib", + "wcs", + ] + + outputs = self.run( + **{key: value for (key, value) in inputs.items() if key in input_keys} + ) + butler_quantum_context.put(outputs, output_refs) From 6f42db427d6824a937ae20a1dc217fc033b93ad7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 29 Apr 2025 21:40:39 +0000 Subject: [PATCH 2/2] Autofix formatting from pre-commit.com hooks --- notebooks/ssi_with_slslim.ipynb | 58 +++++---- slsim/LsstSciencePipeline/generate_catalog.py | 117 +++++++++--------- .../LsstSciencePipeline/inject_slsim_base.py | 6 +- .../LsstSciencePipeline/inject_slsim_visit.py | 2 +- 4 files changed, 100 insertions(+), 83 deletions(-) diff --git a/notebooks/ssi_with_slslim.ipynb b/notebooks/ssi_with_slslim.ipynb index ff3e0f300..dd896dc52 100644 --- a/notebooks/ssi_with_slslim.ipynb +++ b/notebooks/ssi_with_slslim.ipynb @@ -15,6 +15,7 @@ "from astropy.time import Time\n", "import lsst.daf.butler as dafButler\n", "from lsst.pipe.base import Pipeline\n", + "\n", "%matplotlib ipympl" ] }, @@ -26,7 +27,8 @@ "outputs": [], "source": [ "import lsst.afw.display as afwDisplay\n", - "afwDisplay.setDefaultBackend('firefly')\n", + "\n", + "afwDisplay.setDefaultBackend(\"firefly\")\n", "display1 = afwDisplay.Display(frame=1)\n", "display2 = afwDisplay.Display(frame=2)" ] @@ -45,12 +47,12 @@ "sky_area_value = 5.0\n", "\n", "# Butler configuration\n", - "repo_path = \"/repo/main\" \n", + "repo_path = \"/repo/main\"\n", "instrument = \"LSSTComCam\"\n", "\n", "# Specific exposure and detector to process\n", "visit = 2024110800246\n", - "detector = 5 \n", + "detector = 5\n", "\n", "# Collection with pvi images\n", "collection = \"LSSTComCam/runs/DRP/DP1/w_2025_10/DM-49359\"\n", @@ -81,11 +83,13 @@ "metadata": {}, "outputs": [], "source": [ - "gen_cmd = f\"python {slsim_path}/generate_catalog.py --ra-min {ra_min} --ra-max {ra_max} \"\\\n", - " f\" --dec-min {dec_min} --dec-max {dec_max} --start-time {start_time} --end-time {end_time}\"\\\n", - " f\" --n-galaxies {ngal} --sky-area {sky_area} --repo-path {repo_path}\"\\\n", - " f\" --collection {collection} --instrument {instrument}\"\\\n", - " f\" --filters {\" \".join(filters)} --output-file {output_file}\"" + "gen_cmd = (\n", + " f\"python {slsim_path}/generate_catalog.py --ra-min {ra_min} --ra-max {ra_max} \"\n", + " f\" --dec-min {dec_min} --dec-max {dec_max} --start-time {start_time} --end-time {end_time}\"\n", + " f\" --n-galaxies {ngal} --sky-area {sky_area} --repo-path {repo_path}\"\n", + " f\" --collection {collection} --instrument {instrument}\"\n", + " f\" --filters {\" \".join(filters)} --output-file {output_file}\"\n", + ")" ] }, { @@ -95,7 +99,7 @@ "metadata": {}, "outputs": [], "source": [ - "#os.system(gen_cmd)\n", + "# os.system(gen_cmd)\n", "print(gen_cmd)" ] }, @@ -107,9 +111,11 @@ "outputs": [], "source": [ "# Use the ingest_injection_catalog command line tool to ingest the previously created catalog\n", - "injest_cmd = f\"ingest_injection_catalog -b {repo_path} -o {inject_collection}\"\\\n", - " f\" --injection-catalog {output_file} {\" \".join(filters)} -t injection_slsim\"\\\n", - " f\" --format fits\"" + "injest_cmd = (\n", + " f\"ingest_injection_catalog -b {repo_path} -o {inject_collection}\"\n", + " f\" --injection-catalog {output_file} {\" \".join(filters)} -t injection_slsim\"\n", + " f\" --format fits\"\n", + ")" ] }, { @@ -119,7 +125,7 @@ "metadata": {}, "outputs": [], "source": [ - "#os.system(injest_cmd)\n", + "# os.system(injest_cmd)\n", "print(injest_cmd)" ] }, @@ -134,7 +140,8 @@ "pipeline_yaml = \"config_wfake.yaml\"\n", "\n", "with open(pipeline_yaml, \"w\") as f:\n", - " f.write(\"\"\"\n", + " f.write(\n", + " \"\"\"\n", "description: Strong lensing source injection pipeline with DRP processing\n", "instrument: lsst.obs.lsst.LsstComCam\n", "\n", @@ -174,10 +181,13 @@ " - detectAndMeasureDiaSources\n", " description: >\n", " Pipeline for injecting strong lensing sources and running difference imaging\n", - "\"\"\")\n", + "\"\"\"\n", + " )\n", "\n", "print(f\"Created pipeline configuration in {pipeline_yaml}\")\n", - "print(\"The pipeline includes strong lensing injection followed by subtraction and detection tasks, exactly matching test.yaml\")" + "print(\n", + " \"The pipeline includes strong lensing injection followed by subtraction and detection tasks, exactly matching test.yaml\"\n", + ")" ] }, { @@ -187,10 +197,12 @@ "metadata": {}, "outputs": [], "source": [ - "run_cmd = f\"pipetask run -b {repo_path}\"\\\n", - " f\" -d \\\"exposure={visit} and detector={detector}\\\" -p {pipeline_yaml}#injected_DRP\"\\\n", - " f\" -i {collection},{inject_collection}\"\\\n", - " f\" -o {output_collection} --register-dataset-types\"\n" + "run_cmd = (\n", + " f\"pipetask run -b {repo_path}\"\n", + " f' -d \"exposure={visit} and detector={detector}\" -p {pipeline_yaml}#injected_DRP'\n", + " f\" -i {collection},{inject_collection}\"\n", + " f\" -o {output_collection} --register-dataset-types\"\n", + ")" ] }, { @@ -211,7 +223,7 @@ "metadata": {}, "outputs": [], "source": [ - "butler = dafButler.Butler(repo_path,collections=output_collection)" + "butler = dafButler.Butler(repo_path, collections=output_collection)" ] }, { @@ -231,7 +243,9 @@ "metadata": {}, "outputs": [], "source": [ - "diff_pvi = butler.get(\"injected_slsim_goodSeeingDiff_differenceTempExp\", visit=visit, detector=detector)" + "diff_pvi = butler.get(\n", + " \"injected_slsim_goodSeeingDiff_differenceTempExp\", visit=visit, detector=detector\n", + ")" ] }, { diff --git a/slsim/LsstSciencePipeline/generate_catalog.py b/slsim/LsstSciencePipeline/generate_catalog.py index 735f3549e..3da28e8c9 100644 --- a/slsim/LsstSciencePipeline/generate_catalog.py +++ b/slsim/LsstSciencePipeline/generate_catalog.py @@ -19,19 +19,19 @@ HAS_TQDM = False def tqdm(iterable, **kwargs): - """Simple fallback for tqdm progress bar when the library is not available. + """Simple fallback for tqdm progress bar when the library is not + available. :param iterable: Iterable to iterate over :type iterable: iterable - :param kwargs: Keyword arguments that would be passed to tqdm (ignored in fallback) + :param kwargs: Keyword arguments that would be passed to tqdm + (ignored in fallback) :return: Unchanged input iterable :rtype: iterable """ return iterable - - def generate_master_galaxy_list( ra_min, ra_max, dec_min, dec_max, n_galaxies=1000, sky_area_value=0.15 ): @@ -110,7 +110,7 @@ def generate_master_galaxy_list( # Process lens galaxies lens_ell_mask = lens["n_sersic"] < -0.999 n_ell_lens = np.sum(lens_ell_mask) - + if n_ell_lens > 0: phi = np.random.uniform(0, np.pi, size=n_ell_lens) e = lens["ellipticity"][lens_ell_mask].data @@ -135,7 +135,6 @@ def generate_master_galaxy_list( source["e2"][source_ell_mask] = e2 source["n_sersic"][source_ell_mask] = 1 - # Generate random positions within the specified RA/Dec box ra = np.random.uniform(ra_min, ra_max, size=n_galaxies) dec = np.random.uniform(dec_min, dec_max, size=n_galaxies) @@ -199,7 +198,8 @@ def get_calexps_in_region( filters=["u", "g", "r", "i", "z", "y"], max_calexps=10, ): - """Get calexp objects from the butler within given time and spatial constraints. + """Get calexp objects from the butler within given time and spatial + constraints. :param butler: The butler instance :type butler: `lsst.daf.butler.Butler` @@ -330,7 +330,8 @@ def get_calexps_in_region( def find_galaxies_for_each_visit( visit_list, visit_ras, visit_decs, visit_times, galaxy_catalog, reference_time=None ): - """Find which galaxies from the master catalog overlap with each visit's boundaries and record the visit ID and time information. + """Find which galaxies from the master catalog overlap with each visit's + boundaries and record the visit ID and time information. :param visit_list: List of visit dataIds :type visit_list: list @@ -423,83 +424,85 @@ def parse_args(): parser = argparse.ArgumentParser( description="Generate a catalog of galaxies and find overlaps with LSST visits." ) - + # Region of interest - parser.add_argument("--ra-min", type=float, default=52, help="Minimum right ascension in degrees") - parser.add_argument("--ra-max", type=float, default=54, help="Maximum right ascension in degrees") - parser.add_argument("--dec-min", type=float, default=-29, help="Minimum declination in degrees") - parser.add_argument("--dec-max", type=float, default=-27, help="Maximum declination in degrees") - + parser.add_argument( + "--ra-min", type=float, default=52, help="Minimum right ascension in degrees" + ) + parser.add_argument( + "--ra-max", type=float, default=54, help="Maximum right ascension in degrees" + ) + parser.add_argument( + "--dec-min", type=float, default=-29, help="Minimum declination in degrees" + ) + parser.add_argument( + "--dec-max", type=float, default=-27, help="Maximum declination in degrees" + ) + # Time range parser.add_argument( - "--start-time", - type=str, - default="2024-11-01T00:00:00", - help="Start time in ISO format (YYYY-MM-DDTHH:MM:SS)" + "--start-time", + type=str, + default="2024-11-01T00:00:00", + help="Start time in ISO format (YYYY-MM-DDTHH:MM:SS)", ) parser.add_argument( - "--end-time", - type=str, - default="2024-12-29T00:00:00", - help="End time in ISO format (YYYY-MM-DDTHH:MM:SS)" + "--end-time", + type=str, + default="2024-12-29T00:00:00", + help="End time in ISO format (YYYY-MM-DDTHH:MM:SS)", ) - + # Galaxy generation parameters parser.add_argument( - "--n-galaxies", - type=int, - default=10000, - help="Number of galaxies to generate" + "--n-galaxies", type=int, default=10000, help="Number of galaxies to generate" ) parser.add_argument( - "--sky-area", - type=float, - default=5.0, - help="Sky area in square degrees for galaxy simulation" + "--sky-area", + type=float, + default=5.0, + help="Sky area in square degrees for galaxy simulation", ) - + # Butler parameters parser.add_argument( - "--repo-path", - type=str, - default="/repo/main", - help="Path to the butler repository" + "--repo-path", + type=str, + default="/repo/main", + help="Path to the butler repository", ) parser.add_argument( - "--collection", - type=str, - default="LSSTComCam/runs/DRP/DP1/w_2025_07/DM-48940", - help="Butler collection to query" + "--collection", + type=str, + default="LSSTComCam/runs/DRP/DP1/w_2025_07/DM-48940", + help="Butler collection to query", ) parser.add_argument( - "--instrument", - type=str, - default="LSSTComCam", - help="Instrument name to query" + "--instrument", type=str, default="LSSTComCam", help="Instrument name to query" ) parser.add_argument( - "--filters", - type=str, - nargs="+", - default=["r"], - help="List of filters to include" + "--filters", + type=str, + nargs="+", + default=["r"], + help="List of filters to include", ) - + # Output parser.add_argument( - "--output-file", - type=str, - default="galaxy_visit_overlaps.fits", - help="Output file name for galaxy-visit overlaps" + "--output-file", + type=str, + default="galaxy_visit_overlaps.fits", + help="Output file name for galaxy-visit overlaps", ) - + return parser.parse_args() def main(): """Main function to process calexps and find overlapping galaxies.""" args = parse_args() - + # Define region of interest ra_min, ra_max = args.ra_min, args.ra_max dec_min, dec_max = args.dec_min, args.dec_max @@ -512,7 +515,7 @@ def main(): print("Generating master galaxy catalog...") sky_area_value = args.sky_area n_galaxies = args.n_galaxies - + master_galaxies = generate_master_galaxy_list( ra_min=ra_min, ra_max=ra_max, diff --git a/slsim/LsstSciencePipeline/inject_slsim_base.py b/slsim/LsstSciencePipeline/inject_slsim_base.py index 74ffa8923..c988cd6ca 100644 --- a/slsim/LsstSciencePipeline/inject_slsim_base.py +++ b/slsim/LsstSciencePipeline/inject_slsim_base.py @@ -178,7 +178,7 @@ def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs): num_injection_sources = np.sum(good_injections) calib_flux_radius = None num_pix = self.config.stamp_size - + if num_injection_sources > 0: cosmo = FlatLambdaCDM(H0=70, Om0=0.3) @@ -288,14 +288,14 @@ def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs): source_redshift=source.redshift, deflector_redshift=lens.redshift, ) - + lens_class = Lens( source_class=source, deflector_class=lens, cosmo=cosmo, los_class=los, ) - + lens_im = lens_image( lens_class=lens_class, band=band, diff --git a/slsim/LsstSciencePipeline/inject_slsim_visit.py b/slsim/LsstSciencePipeline/inject_slsim_visit.py index 5b9884d47..f5bb8ad64 100644 --- a/slsim/LsstSciencePipeline/inject_slsim_visit.py +++ b/slsim/LsstSciencePipeline/inject_slsim_visit.py @@ -164,7 +164,7 @@ def runQuantum(self, butler_quantum_context, input_refs, output_refs): "photo_calib", "wcs", ] - + outputs = self.run( **{key: value for (key, value) in inputs.items() if key in input_keys} )