
Generating a sample-based line-of-sight redshift prior
======================================================

The sample-based LOS redshift prior is an alternative to the grid-based approach.
Instead of evaluating $p(z|\Omega)$ on a fixed redshift grid, it draws
$(z, M, w)$ samples per HEALPix pixel using importance sampling and stores them
in a [zarr](https://zarr.dev/) directory store.

**Advantages over the grid-based approach:**

- Galaxy weighting is **not baked in** at generation time. Weights are stored
  unweighted; consumers apply any weighting scheme (luminosity, uniform, SFR, etc.)
  downstream without regenerating the prior.
- Multi-nside support via **child-pixel aggregation**: generate at the finest
  nside needed, then combine child pixels for coarser resolutions at read time.
- **Parallel-safe output**: zarr's directory-per-array layout allows multiple
  workers to write to different pixels simultaneously without a writer thread.


## Quick start

```bash
gwcosmo_compute_redshift_prior_samples \
    --catalog GLADE+ \
    --catalog_band K \
    --nside 64 \
    --coarse_nside 8 \
    --H0 70 --Omega_m 0.3065 --w0 -1 --wa 0 \
    --zmax 10 \
    --Kcorrections True \
    --nsamps_per_galaxy 100 \
    --nsamps_ooc 100000 \
    --num_threads 8 \
    --output_zarr my_los_prior.zarr
```

This produces a zarr store at `my_los_prior.zarr` containing
out-of-catalogue (OOC) samples (generated once globally) and
in-catalogue (IC) samples per HEALPix pixel.


## Command-line arguments

All arguments are shared with the rest of gwcosmo via the central argument
parser in `gwcosmo.utilities.arguments`. The sample-specific arguments are:

| Argument | Default | Description |
|---|---|---|
| `--nsamps_per_galaxy` | 100 | Number of redshift samples drawn per catalogue galaxy |
| `--nsamps_ooc` | 100000 | Number of out-of-catalogue samples (cosmological prior $\times$ Schechter function) |
| `--output_zarr` | auto | Path for the output `.zarr` store. If not given, a name is generated from the catalogue, band, nside, and pixel index |

Standard arguments such as `--catalog`, `--catalog_band`, `--nside`,
`--Kcorrections`, `--zmax`, `--zcut`, `--mth`, `--H0`, `--Omega_m`, `--w0`,
`--wa`, `--num_threads`, `--schech_alpha`, `--schech_Mstar`, `--schech_Mmin`,
`--schech_Mmax`, `--min_gals_for_threshold`, `--maps_path`, `--coarse_nside`,
and `--pixel_index` are documented in `gwcosmo_compute_redshift_prior_samples --help`.


## Zarr store layout

```
my_los_prior.zarr/
  .zattrs                  # root metadata (see below)
  ooc/
    zsamps                 # (N,) float32 — redshift samples
    Msamps                 # (N,) float32 — absolute magnitude samples
    weights                # (N,) float32 — importance weights
  pixels/
    {pixel_index}/
      .zattrs              # {mth, zcut, ngals, empty_pixel}
      zsamps               # (M,) float32 — IC redshift samples
      Msamps               # (M,) float32 — IC absolute magnitude samples
      weights              # (M,) float32 — IC importance weights
      apparent_mags        # (ngals,) float32 — one per galaxy
```

**Root metadata** (`.zattrs`):

| Key | Description |
|---|---|
| `H0`, `Omega_m`, `w0`, `wa` | Cosmological parameters used during generation |
| `zmax` | Maximum redshift |
| `catalog`, `band` | Galaxy catalogue name and photometric band |
| `nside` | HEALPix resolution |
| `Kcorrections` | Whether K-corrections were applied |
| `schech_alpha`, `schech_Mstar`, `schech_Mmin`, `schech_Mmax` | Schechter function parameters |
| `nsamps_per_galaxy`, `nsamps_ooc` | Sample counts |
| `gwcosmo_version` | Version of gwcosmo used |
| `created_utc` | Timestamp |
| `weighting` | Always `"unweighted"` — weighting is applied downstream |

**Per-pixel metadata** (`pixels/{index}/.zattrs`):

| Key | Description |
|---|---|
| `empty_pixel` | `true` if no galaxies pass selection; no arrays are stored |
| `mth` | Apparent magnitude threshold for this pixel |
| `zcut` | Redshift cut applied |
| `ngals` | Number of galaxies in the pixel |

Empty pixels (those with `mth = inf` or too few galaxies) store only
metadata with `empty_pixel: true` and no sample arrays.


## Sampling strategy

### Out-of-catalogue (OOC) samples

OOC samples represent the population of galaxies **not** in the catalogue.
They are pixel-independent and stored once globally.

- **Redshift**: drawn via log-space importance sampling from the cosmological
  prior $p(z)$. Sampling uniformly in $\log z$ and weighting by $p(z) \cdot z$
  provides better coverage at low redshifts.
- **Absolute magnitude**: drawn via inverse-transform sampling from the
  Schechter magnitude function $\phi(M | H_0)$.

### In-catalogue (IC) samples

IC samples are generated per HEALPix pixel from the galaxies in the catalogue.

For each galaxy $i$ in the pixel with observed redshift $z_i$ and uncertainty
$\sigma_{z,i}$:

1. Draw `nsamps_per_galaxy` redshift samples from
   $\mathcal{N}_{\mathrm{trunc}}(z_i, \sigma_{z,i})$ (truncated at $z > 0$).
2. Compute the absolute magnitude $M$ from the galaxy's apparent magnitude $m$
   and luminosity distance $d_L(z)$, applying K-corrections if enabled.
3. Assign uniform weights $1 / n_{\mathrm{samps}}$ per galaxy.

The apparent magnitude of each galaxy is also stored to support
recomputing $m_{\mathrm{th}}$ when aggregating to coarser resolutions.


## Reading the zarr store in Python

The `gwcosmo.prior.catalog_samples_io` module provides convenience functions
for loading and combining the stored samples.

### Opening the store

```python
from gwcosmo.prior.catalog_samples_io import open_los_samples

store = open_los_samples("my_los_prior.zarr")
print(dict(store.attrs))  # root metadata
```

### Loading OOC samples

```python
from gwcosmo.prior.catalog_samples_io import load_ooc_samples

zsamps, Msamps, weights = load_ooc_samples(store)
```

### Loading a single pixel

```python
from gwcosmo.prior.catalog_samples_io import load_pixel_samples

pix = load_pixel_samples(store, pixel_index=42)
if pix['empty_pixel']:
    print("No catalogue support in this pixel")
else:
    print(f"ngals={pix['ngals']}, mth={pix['mth']:.2f}")
    z = pix['zsamps']
    M = pix['Msamps']
    w = pix['weights']
```

### Combining IC + OOC with EM selection

The `load_pixel_los_prior` function applies the electromagnetic selection
criteria (apparent magnitude threshold and redshift cut) and combines the
in-catalogue and out-of-catalogue contributions:

```python
from gwcosmo.prior.catalog_samples_io import load_pixel_los_prior

zsamps, Msamps, weights = load_pixel_los_prior(store, pixel_index=42)
```

An optional `weighting_func` can be passed to apply galaxy weighting:

```python
from gwcosmo.utilities.host_galaxy_merger_relations import LuminosityWeighting

lw = LuminosityWeighting()
zsamps, Msamps, weights = load_pixel_los_prior(
    store, pixel_index=42, weighting_func=lw
)
```

### Multi-nside aggregation

If the store was generated at a fine nside (e.g. 64), you can obtain
results at a coarser nside (e.g. 32) by aggregating child pixels.
This concatenates IC samples from all child pixels and recomputes
$m_{\mathrm{th}}$ as the median of the combined apparent magnitudes:

```python
from gwcosmo.prior.catalog_samples_io import aggregate_child_pixels

result = aggregate_child_pixels(store, parent_nside=32, parent_pixel_index=0)
if result is not None:
    print(f"Aggregated: {result['ngals']} galaxies, mth={result['mth']:.2f}")
```

This uses nested HEALPix ordering: for a store at `nside_fine`, the
child pixels of a coarser parent pixel are a contiguous index range.


## Relationship to the grid-based prior

The grid-based approach (`gwcosmo_compute_redshift_prior`) evaluates
$p(z|\Omega)$ on a fixed redshift grid and stores the result in HDF5.
Galaxy weighting is applied at generation time. This is the established
production method and remains the default for standard analyses.

The sample-based approach is designed for cases where:

- You need to experiment with different galaxy weighting schemes
  without regenerating the prior each time.
- You want to use results at multiple HEALPix resolutions from a
  single generation run.
- You want to inspect or post-process the individual galaxy
  contributions to the prior.


## API reference

Detailed API documentation is available in the {doc}`API Reference <api/prior>`.
The key modules are:

- {py:mod}`gwcosmo.prior.LOS_redshift_prior_samples` — sampling classes
  (`LineOfSightRedshiftPriorOutofCatalogue`,
  `LineOfSightRedshiftPriorInCatalogue`,
  `LineOfSightRedshiftPriorSinglePixel`)
- {py:mod}`gwcosmo.prior.catalog_samples_io` — reader utilities
  (`open_los_samples`, `load_ooc_samples`, `load_pixel_samples`,
  `load_pixel_los_prior`, `aggregate_child_pixels`)
