From bd17081332b424dd4d4f752ec0e801ad258f7fd4 Mon Sep 17 00:00:00 2001
From: Michael Troxel <mat90@dcc-login-04.oit.duke.edu>
Date: Thu, 11 Apr 2024 13:56:18 -0400
Subject: [PATCH] Addition of sampling over nz modes

---
 modules_desy6/photoz_uz/module.yaml  | 51 ++++++++++++++++++++++++++++
 modules_desy6/photoz_uz/photoz_uz.py | 36 ++++++++++++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 modules_desy6/photoz_uz/module.yaml
 create mode 100644 modules_desy6/photoz_uz/photoz_uz.py

diff --git a/modules_desy6/photoz_uz/module.yaml b/modules_desy6/photoz_uz/module.yaml
new file mode 100644
index 00000000..f1c09418
--- /dev/null
+++ b/modules_desy6/photoz_uz/module.yaml
@@ -0,0 +1,51 @@
+#This is a template for module description files
+name: photoz_uz
+version: 1
+purpose: Modify a set of loaded n(z) distributions by a set of eignmodes and amplitudes.
+url: ''
+interface: photoz_uz.py
+attribution: Troxel & Gary Bernstein
+rules: ''
+cite: []
+assumptions:
+-   files containing information on n(z) modes and priors on amplitudes
+
+explanation: |
+    
+params:
+    sample:
+        meaning: name of sample
+        type: str
+        default: ""
+    basis_file:
+        meaning: file containing n(z) mode array
+	type: str
+	default: ""
+    n_modes:
+        meaning: number of modes to use
+        type: int
+        default: 0
+inputs:
+    wl_number_density:
+        nbin:
+            meaning: Number of redshift bins
+            type: int
+            default:
+        z:
+            meaning: Redshift sample points of n(z) estimates
+            type: real 1d
+            default:
+        bin_i:
+            meaning: n(z)for i=1..nbin. n(z) estimates
+            type: real 1d
+            default:
+    {sample}_photoz_u:
+        bias_i:
+            meaning: For i=1..nmode. Amplitude for mode i.
+            type: real
+            default:
+outputs:
+    wl_number_density:
+        bin_i:
+            meaning: n(z) for i=1..nbin. Modified n(z) estimates replaced old value
+            type: real 1d
diff --git a/modules_desy6/photoz_uz/photoz_uz.py b/modules_desy6/photoz_uz/photoz_uz.py
new file mode 100644
index 00000000..64aca26d
--- /dev/null
+++ b/modules_desy6/photoz_uz/photoz_uz.py
@@ -0,0 +1,36 @@
+from builtins import range
+from cosmosis.datablock import option_section, names
+import numpy as np
+
+def setup(options):
+    sample = options.get_string(option_section, "sample", "")
+    bias_section = sample+"_photoz_u"
+    basis_file = options.get_string(option_section, "basis_file", "")
+    n_modes = options.get_int(option_section, "n_modes", 0)
+    U = np.loadtxt(basis_file)[:,:n_modes]
+    return {"sample": sample,
+            "bias_section": bias_section,
+            "basis":U,
+            "n_modes":n_modes}
+
+def execute(block, config):
+    pz = 'nz_'+config['sample']
+    uvals = config['bias_section']
+    U = config['basis']
+    n_modes = config['n_modes']
+    nbin = block[pz, "nbin"]
+    if nbin>4:
+        nbin=4
+    z = block[pz, "z"]
+    dz = U@np.array([block[uvals, "u_%d" % j] for j in range(n_modes)])
+    dz = dz.reshape((nbin,len(dz)//nbin))
+    for i in range(1,nbin+1):
+        bin_name = "bin_%d" % i
+        nz = block[pz, bin_name]
+        nz[1:]+=dz[i-1,:]
+        nz /= np.trapz(nz, z)
+        block[pz, bin_name] = nz
+    return 0
+
+def cleanup(config):
+    pass