Tqz/somoclu pipe

nb/estimation/som_train_estimate.ipynb

@@ -0,0 +1,264 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f9463120",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# usual imports\n",
+    "import os\n",
+    "import numpy as np\n",
+    "from rail.utils.path_utils import find_rail_file\n",
+    "from rail.pipelines.estimation.somoclu_train_estimate import SOMTrainEstimatePipeline\n",
+    "from rail.core import common_params\n",
+    "import ceci"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "916a2494",
+   "metadata": {},
+   "source": [
+    "### Set common parameters for the photometric catalog\n",
+    "\n",
+    "We use the small DC2-like test catalogs bundled with RAIL.\n",
+    "Their magnitude columns follow the `mag_{band}_lsst` / `mag_err_{band}_lsst`\n",
+    "naming convention, with a `photometry` HDF5 group and a `redshift` column."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3cc64157",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bands = 'ugrizy'\n",
+    "band_cols     = [f'mag_{b}_lsst'     for b in bands]\n",
+    "err_band_cols = [f'mag_err_{b}_lsst' for b in bands]\n",
+    "\n",
+    "maglim_dict = {\n",
+    "    'mag_u_lsst': 24.0,\n",
+    "    'mag_g_lsst': 27.66,\n",
+    "    'mag_r_lsst': 27.25,\n",
+    "    'mag_i_lsst': 26.6,\n",
+    "    'mag_z_lsst': 26.24,\n",
+    "    'mag_y_lsst': 25.35,\n",
+    "}\n",
+    "\n",
+    "common_params.set_param_defaults(\n",
+    "    bands=band_cols,\n",
+    "    err_bands=err_band_cols,\n",
+    "    nondetect_val=np.nan,\n",
+    "    ref_band='mag_i_lsst',\n",
+    "    redshift_col='redshift',\n",
+    "    mag_limits=maglim_dict,\n",
+    "    zmax=3.0,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "090457b4",
+   "metadata": {},
+   "source": [
+    "### Configure the SOM informer and summarizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d966ae99",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inform_dict = dict(\n",
+    "    hdf5_groupname='photometry',\n",
+    "    n_rows=10,\n",
+    "    n_columns=10,\n",
+    "    gridtype='hexagonal',\n",
+    "    std_coeff=1.0,\n",
+    "    som_learning_rate=0.3,\n",
+    "    n_epochs=2,\n",
+    "    ref_column_name='mag_i_lsst',\n",
+    "    initialization='random',\n",
+    "    column_usage='magandcolors',\n",
+    ")\n",
+    "\n",
+    "summ_dict = dict(\n",
+    "    hdf5_groupname='photometry',\n",
+    "    spec_groupname='photometry',\n",
+    "    nzbins=301,\n",
+    "    nsamples=20,\n",
+    "    objid_name='id',\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "19067b35",
+   "metadata": {},
+   "source": [
+    "### Build the pipeline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "36ab51b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pipe = SOMTrainEstimatePipeline(inform_dict=inform_dict, summ_dict=summ_dict)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "96ecd511",
+   "metadata": {},
+   "source": [
+    "### Locate the input data files\n",
+    "\n",
+    "Three separate catalogs are used:\n",
+    "\n",
+    "| Role | File | Stage |\n",
+    "|------|------|-------|\n",
+    "| `input_train` | `training_100gal.hdf5` | `inform_som` — trains the SOM |\n",
+    "| `input_photo` | `validation_10gal.hdf5` | `summarize_som` — photometric catalog whose n(z) we want |\n",
+    "| `spec_input`  | `training_100gal.hdf5` | `summarize_som` — spectroscopic reference with true redshifts |"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15567a43",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_file = find_rail_file('examples_data/testdata/training_100gal.hdf5')\n",
+    "photo_file = find_rail_file('examples_data/testdata/validation_10gal.hdf5')\n",
+    "spec_file  = find_rail_file('examples_data/testdata/training_100gal.hdf5')\n",
+    "\n",
+    "output_dir = os.path.join('projects', 'som_test')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "29ebaff1",
+   "metadata": {},
+   "source": [
+    "### Provide input files and initialise the pipeline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4a5afa9c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_dict = pipe.default_input_dict.copy()\n",
+    "input_dict.update(\n",
+    "    input_train=train_file,\n",
+    "    input_photo=photo_file,\n",
+    "    spec_input=spec_file,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "088b42d7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pipe_info = pipe.initialize(input_dict, dict(output_dir=output_dir, log_dir='.', resume=True), None)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f83674dd-2779-445a-a457-c62db252b8c5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pipe.print_stages()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b7b1ebdf",
+   "metadata": {},
+   "source": [
+    "### Save the pipeline to a YAML file"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "298e80fa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pipe.save('som_train_estimate.yml')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d6013957",
+   "metadata": {},
+   "source": [
+    "[For NERSC / cluster users!]\n",
+    "\n",
+    "This won't work on a login node / Jupyter server without compute access. To run\n",
+    "the pipeline in batch, you need to:\n",
+    "1. Add `name: local` to the `site` section in `som_train_estimate.yml`.\n",
+    "2. SSH into a compute node, activate the RAIL environment, and run\n",
+    "   `ceci som_train_estimate.yml`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "97dd95f0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pr = ceci.Pipeline.read('som_train_estimate.yml')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d8a21466",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pr.run()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "rail_env_dev",
+   "language": "python",
+   "name": "rail_env_dev"
+  },
+  "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.14.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/rail/pipelines/estimation/somoclu_train_estimate.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+from rail.core.stage import RailStage, RailPipeline
+from rail.estimation.algos.somoclu_som import SOMocluInformer, SOMocluSummarizer
+
+
+class SOMTrainEstimatePipeline(RailPipeline):
+    """Pipeline that trains a SOMoclu SOM and uses it to estimate n(z).
+
+    Stages
+    ------
+    inform_som : SOMocluInformer
+        Trains the self-organising map on the photometric input catalog.
+    summarize_som : SOMocluSummarizer
+        Uses the trained SOM together with a spectroscopic reference sample
+        to produce an n(z) ensemble.
+    """
+
+    default_input_dict = {
+        'input_train': 'dummy.in',
+        'input_photo': 'dummy.in',
+        'spec_input': 'dummy.in',
+    }
+
+    def __init__(self, inform_dict=None, summ_dict=None):
+        RailPipeline.__init__(self)
+
+        if inform_dict is None:
+            inform_dict = {}
+        if summ_dict is None:
+            summ_dict = {}
+
+        informer = SOMocluInformer.make_and_connect(
+            name='inform_som',
+            aliases=dict(input='input_train'),
+            **inform_dict,
+        )
+        self.add_stage(informer)
+
+        summarizer = SOMocluSummarizer.make_and_connect(
+            name='summarize_som',
+            aliases=dict(input='input_photo'),
+            connections=dict(model=informer.io.model),
+            **summ_dict,
+        )
+        self.add_stage(summarizer)

tests/test_pipelines.py

@@ -15,6 +15,7 @@
         'rail.pipelines.utils.prepare_observed.PrepareObservedPipeline',
         'rail.pipelines.examples.goldenspike.goldenspike.GoldenspikePipeline',
         'rail.pipelines.examples.survey_nonuniformity.survey_nonuniformity.SurveyNonuniformDegraderPipeline',
+        'rail.pipelines.estimation.somoclu_train_estimate.SOMTrainEstimatePipeline',
     ]
 )
 def test_build_and_read_pipeline(pipeline_class):