Pipeline

mejiro_pipeline

Step 0: Cache PSFs

Generate and cache Roman PSFs in parallel.

This script reads a YAML configuration file specifying parameters for PSF generation, such as oversampling factors, bands, detectors, detector positions, and pixel sizes. It determines which PSFs need to be generated based on existing cache files, and spins up multiple processes to compute and save the missing PSFs as .npy files.

Usage:

python3 _00_cache_psfs.py –config <config.yaml> [–data_dir <output_dir>]

Arguments:

–config: Path to the YAML configuration file. –data_dir: Optional override for the data directory specified in the config file.

Step 1a: Generate Galaxy Tables

Pre-generates galaxy population tables for the survey simulation.

This script generates a configurable number of galaxy tables by running SkyPy and SLHammocks pipelines, then serializes them as pickle files. This separates the expensive galaxy population generation from the survey simulation itself, allowing the simulation step (_01b) to load pre-computed tables and skip the initialization.

Each table is generated with a unique random seed for reproducibility. Tables are associated with detectors (for Roman, each table is assigned a detector via round-robin), since SkyPy configs differ per detector.

Usage:

python3 _01a_generate_galaxy_tables.py –config <config.yaml> [–data_dir <output_dir>]

Arguments:

–config: Path to the YAML configuration file. –data_dir: Optional override for the data directory specified in the config file.

Step 1b: Run Survey Simulation

Simulates a strong lensing survey using pre-computed galaxy tables from _01a.

Loads galaxy population tables generated by _01a_generate_galaxy_tables.py, reconstructs the lens population objects, and runs the survey simulation. This avoids the expensive SkyPy/SLHammocks initialization that dominated runtime in the original single-step design. Each run is assigned a galaxy table via round-robin (run_index % num_galaxy_tables), with per-run random seeding for reproducible, unique draws across runs.

When jaxtronomy.use_jax is True, lenstronomy ray-shooting (invoked during per-candidate SNR computation) is routed through jaxtronomy by flipping StrongLens.use_jax to True on every profile in _JAXXED_MODELS after building each GalaxyGalaxy via from_slsim. When False, ray-shooting stays in pure lenstronomy and workers never import JAX (the spawn pool, BLAS threading controls, and CPU affinity pinning below apply either way).

With JAX enabled, execution topology is selected from the YAML config’s jaxtronomy.jax_platform:

• cpu – ProcessPoolExecutor with a spawn multiprocessing context. Each worker imports JAX after JAX_PLATFORM_NAME is set, and all workers share an on-disk JAX compilation cache so the JIT cost for each unique lens_model_list signature is paid only once across the whole run.

• gpu – single-process sequential execution over runs. One in-flight kernel per GPU; oversubscribing N processes onto one device thrashes memory.

Worker threading (CPU path):

Three layered controls keep N workers from collectively oversubscribing the host:

1. BLAS/OMP env vars (OMP_NUM_THREADS etc.) set below before numpy imports, so NumPy/SciPy/galsim stay single-threaded in every worker.

2. XLA_FLAGS set in _worker_init when JAX is enabled (unconditional assignment, not setdefault, so an inherited shell value cannot mask it): --xla_cpu_multi_thread_eigen=false disables Eigen op-level parallelism; --xla_force_host_platform_device_count=1 keeps JAX to one logical CPU device per worker.

3. os.sched_setaffinity in _worker_init pins each worker to a single core (ids handed out via a multiprocessing.Queue built in _run_spawn_pool). This is the actual hard cap: JAX’s CPU PJRT client still allocates dispatch and JIT-compile threads sized to hardware_concurrency() that XLA_FLAGS does not bound, but affinity restricts them to one core so the system run queue stays near workers rather than workers * ~3.

Usage:

python3 _01b_run_survey_simulation.py –config <config.yaml> [–data_dir <output_dir>] [–resume]

Step 2: Build Lens List

Builds a list of StrongLens objects from previously detected lensing systems.

This script processes the output of a strong lensing survey simulation, converting detected lensing systems into mejiro StrongLens objects. It reads configuration parameters from a mejiro configuration YAML file and supports multiple instruments (Roman, HWO). Output lenses are pickled for downstream analysis.

Usage:

python3 _02_build_lens_list.py –config <config.yaml> [–data_dir <output_dir>]

Arguments:

–config: Path to the YAML configuration file. –data_dir: Optional override for the data directory specified in the config file.

Step 3: Generate Subhalo Realizations

Step 4: Create Synthetic Images (Ray-Shooting)

Generates synthetic images: idealized images with no noise or detector effects (optionally, convolved with PSF).

The YAML config’s jaxtronomy.use_jax selects the execution path:

• use_jax: False – lenstronomy ray-shooting in a fork-context ProcessPoolExecutor.

• use_jax: True – ray-shooting routed through jaxtronomy by flipping StrongLens.use_jax to True on every profile in _JAXXED_MODELS before constructing the SyntheticImage. Execution topology is selected from jaxtronomy.jax_platform:

  • cpu – ProcessPoolExecutor with a spawn multiprocessing context. Each worker imports JAX after JAX_PLATFORM_NAME is set, and all workers share an on-disk JAX compilation cache so the JIT cost for each unique lens_model_list signature is paid only once across the whole run.

  • gpu – single-process sequential execution, bucketed by lens model signature so the JIT cache is reused within each bucket. One in-flight kernel per GPU; oversubscribing N processes onto one device thrashes memory.

Worker threading (JAX CPU path):

Three layered controls keep N workers from collectively oversubscribing the host:

1. BLAS/OMP env vars (OMP_NUM_THREADS etc.) set below before numpy imports, so NumPy/SciPy/galsim stay single-threaded in every worker.

2. XLA_FLAGS set in _jax_worker_init (unconditional assignment, not setdefault, so an inherited shell value cannot mask it): --xla_cpu_multi_thread_eigen=false disables Eigen op-level parallelism; --xla_force_host_platform_device_count=1 keeps JAX to one logical CPU device per worker.

3. os.sched_setaffinity in _jax_worker_init pins each worker to a single core (ids handed out via a multiprocessing.Queue built in _run_jax_cpu). This is the actual hard cap: JAX’s CPU PJRT client still allocates dispatch and JIT-compile threads sized to hardware_concurrency() that XLA_FLAGS does not bound, but affinity restricts them to one core so the system run queue stays near workers rather than workers * ~3.

Input comes from step 02 or 03 (--prev-step, default 03); outputs land in <pipeline_dir>/04/ for consumption by the 05 scripts.

Usage:

python3 _04_create_synthetic_images.py –config <config.yaml> [–prev-step {02,03}] [–resume] [–sequential]

Step 5: Create Exposures (galsim)

Generates exposures from synthetic images, i.e., apply sky background and detector effects to idealized images.

This script processes synthetic images produced in previous pipeline steps, generating exposures for each lensing system using instrument-specific parameters and simulation engines. It reads a mejiro YAML configuration file specifying exposure options. Multiprocessing is used to parallelize exposure creation across available CPU cores.

Usage:

python3 _05_galsim.py –config <config.yaml> [–data_dir <output_dir>]

Arguments:

–config: Path to the YAML configuration file. –data_dir: Optional override for the data directory specified in the config file.

Step 5 (Alternative): Create Exposures (romanisim)

Runs romanisim detector simulation on tiled synthetic images, producing either L2 (single-exposure) or L3 (dithered, co-added mosaic) cutouts.

--level l2 (default) tiles SyntheticImage pickles into one 4088x4088 detector array, runs romanisim once to apply detector effects, and extracts individual cutouts. The tile size is derived from config[‘synthetic_image’][‘fov_arcsec’] (via util.set_odd_num_pix at Roman’s 0.11”/pix), and the grid is built dynamically as grid_side = 4088 // tile_size tiles per side, origin-aligned to (0,0) and leaving the right/top remainder empty. Systems are processed in batches of grid_side**2 until all systems for each SCA/band are complete. Output lands in <data_dir>/<pipeline_label>/05_romanisim/sca##/ in L2 DN/s.

--level l3 instead:

1. Tiles SyntheticImages only into the region where all dither pointings of --dither-pattern overlap on the detector. For a gap-filling pattern like BOXGAP4_1 (~205 arcsec offsets) this is the ~one-quadrant 4-fold-overlap region; for a sub-pixel pattern like SUB4 it is essentially the entire detector, so each batch carries ~4x more systems at ~4x less compute for the same effective depth.

2. Runs romanisim once per dither pointing to make one L2 exposure each.

3. Co-adds the L2s with romancal’s MosaicPipeline into one L3 mosaic. Every system is covered by all exposures, so the mosaic has uniform weight and an effective exposure time of n_dithers x the single-exposure time.

4. Extracts each system’s cutout from the mosaic (via the mosaic WCS at the system’s sky position) and saves it in the exact same output format as --level l2, except the pixel values are the mosaic’s native MJy/sr surface brightness rather than L2 DN/s.

Every invocation writes to 05_romanisim/ regardless of --level or --dither-pattern: the step directory names the step, not the flags it was run with. Trade-off of sub-pixel patterns (those in WfiImagingSubpixel.txt, e.g. SUB4): every exposure puts a system on nearly the same pixels, so fixed-pattern detector effects (flat error, hot pixels) stay correlated across the stack instead of averaging down the way they do with large dithers. Per-exposure noise still decorrelates (each dither gets its own rng seed).

Multiprocessing is used to parallelize batch processing.

Usage:

python3 _05_romanisim.py –config <config.yaml> [–data_dir <dir>] [–level {l2,l3}]

[–dither-pattern <name>] [–sequential] [–max-systems N] [–resume]

Arguments:

–config: Path to the YAML configuration file. –data_dir: Parent directory of pipeline output. Overrides data_dir in the config file. –level: ‘l2’ (default) for single-exposure cutouts, ‘l3’ for co-added mosaic cutouts. –dither-pattern: Dither pattern name for –level l3 (default BOXGAP4_1); sub-pixel

patterns from WfiImagingSubpixel.txt (e.g. SUB4) are also accepted. Note the number of dither steps sets the effective depth.

–sequential: Process systems sequentially from the start instead of randomly when a

limit is imposed.

–max-systems: Process only the first N systems (lowest UIDs) across all SCAs, in

total (unlike config[‘limit’], which is applied per SCA).

–resume: Preserve existing output and skip already-completed batches (those with a

batch_complete_*.txt sentinel). Default is to delete existing output and rebuild from scratch. Note the sentinels record neither –level nor –dither-pattern, and all variants share the one 05_romanisim/ directory, so only resume with the same –level and –dither-pattern the original run used – otherwise the resumed run will mix incompatible products. The default (non—resume) wipe is what keeps a run with different flags from silently mixing with an earlier one.

mejiro.pipeline._05_romanisim.compute_overlap_skygrid(pointings, sca, band, ma_table_number, read_pattern, date, tile_size)

Build the grid of source sky positions that land on-detector in every dither.

Mirrors mejiro.point_wfi.find_source_position() but, instead of returning a single position, lays a regular pitch-spaced grid of tile_size x tile_size tile centers across the rectangular intersection (in dither-0 pixel coordinates) of all dither footprints. pitch is tile_size + DISTORTION_GUARD.

Returns a tuple (wcses, source_skies) where wcses are the per-dither galsim WCSes and source_skies is a list of galsim.CelestialCoord tile-center positions (row-major).

mejiro.pipeline._05_romanisim.exposure_cutout_name(input_path)

Derive the Exposure .npy cutout filename from a SyntheticImage input path.

Handles both full .pkl pickles and lightweight .npz inputs by stripping whatever extension is present before appending .npy (avoids .npz.npy).

Calculate SNRs

Calculates signal-to-noise ratios (SNRs) for simulated exposures.

This script computes SNR values for each lensing system processed in previous pipeline steps and saves name-SNR pairs for downstream filtering or analysis. It reads a YAML configuration file specifying SNR calculation parameters and supports both sequential and parallel processing modes.

Usage:

python3 calculate_snrs.py –config <config.yaml> [–data_dir <dir>] [–sequential] [–resume] [–prev-step <step_dir>]

Arguments:

–config: Path to the YAML configuration file. –data_dir: Parent directory of pipeline output. Overrides data_dir in the config file. –sequential: Run in sequential mode instead of parallel. –resume: Preserve existing output and skip already-completed items. Default is to delete and rebuild from scratch. –prev-step: Pipeline step directory holding the input exposures (default: ‘05_romanisim’).

Step 6: Export Dataset

Exports romanisim exposures and synthetic images to HDF5 format.

This script reads romanisim exposure cutouts (.npy) and the corresponding SyntheticImage files from previous pipeline steps, and writes them to an HDF5 file with relevant metadata:

• exposure data: loaded from .npy cutout files (05_romanisim/sca*/)

• exposure_time: from config[‘exposure’][‘ma_table_number’] via romanisim parameters

• lens and synthetic image metadata: from the SyntheticImage files in step 04

• SNR: read from name_snr_pairs.pkl produced by calculate_snrs.py (optional)

• units: DN/s

If config[‘dataset’][‘labeled’] is False, truth attributes (main halo mass, Einstein radius, velocity dispersion, substructure) are omitted from the HDF5 file and an answer-key CSV is written alongside it.

Usage:

python3 _06_h5_export.py –config <config.yaml> [–data_dir <dir>] [–prev-step <step_dir>]

Arguments:

–config: Path to the YAML configuration file. –data_dir: Parent directory of pipeline output. Overrides data_dir in the config file. –prev-step: Pipeline step directory holding the input exposures (default: ‘05_romanisim’).