Add multiple FeH components to GaussianFeH model

src/common/otable.cpp

@@ -528,6 +528,7 @@ int otable::getSortedColumnsForOutput(std::vector<const columndef*> &outColumns)
 	field["absmag"] = ord++;
 	field["comp"] = ord++;
 	field["FeH"] = ord++;
+	field["FeH_comp"] = ord++;
 	field["vcyl"] = ord++;
 	field["pmlb"] = ord++;
 	field["pmradec"] = ord++;

src/modules/GaussianFeH_gpu.cu.h

@@ -21,6 +21,19 @@
 #ifndef GaussianFeH_gpu_cu_h__
 #define GaussianFeH_gpu_cu_h__
 
+//
+// Module data passed to device kernel
+//
+struct os_GaussianFeH_data
+{
+	static const int MAXCOMP = 10;
+
+	int ncomp;
+	float mean[MAXCOMP];
+	float sigma[MAXCOMP];
+	float pcum[MAXCOMP];
+};
+
 //
 // Device kernel implementation
 //
@@ -29,12 +42,13 @@
 KERNEL(
 	ks, 3*4,
 	os_GaussianFeH_kernel(
-		otable_ks ks, bit_map applyToComponents, float mean, float sigma, gpu_rng_t rng,
+		otable_ks ks, bit_map applyToComponents, os_GaussianFeH_data par, gpu_rng_t rng,
 		cint_t::gpu_t comp, cint_t::gpu_t hidden,
 		cfloat_t::gpu_t XYZ,
+		cint_t::gpu_t FeH_comp,
 		cfloat_t::gpu_t FeH),
 	os_GaussianFeH_kernel,
-	(ks, applyToComponents, mean, sigma, rng, comp, hidden, XYZ, FeH)
+	(ks, applyToComponents, par, rng, comp, hidden, XYZ, FeH_comp, FeH)
 )
 {
 	uint32_t tid = threadID();
@@ -46,7 +60,14 @@ KERNEL(
 		int cmp = comp(row);
 		if(!applyToComponents.isset(cmp)) { continue; }
 
-		FeH(row) = mean + rng.gaussian(sigma);
+		int at = 0;
+		if(par.ncomp != 1)
+		{
+			float u = rng.uniform();
+			while(par.pcum[at] < u) ++at;
+		}
+		FeH(row) = par.mean[at] + rng.gaussian(par.sigma[at]);
+		FeH_comp(row) = at;
 	}
 	rng.store(tid);
 }

src/modules/GaussianFeH_host.cpp

@@ -24,14 +24,16 @@
 #include "../pipeline.h"
 #include "GaussianFeH_gpu.cu.h"
 
+#include <numeric>
+
 #include <astro/system/config.h>
 #include <astro/useall.h>
 
 // os_GaussianFeH -- Generate Fe/H based on Ivezic et al (2008)
 class os_GaussianFeH : public osink
 {
 protected:
-	float mean, sigma;
+	os_GaussianFeH_data par;
 
 public:
 	virtual size_t process(otable &in, size_t begin, size_t end, rng_t &rng);
@@ -42,15 +44,17 @@ class os_GaussianFeH : public osink
 	os_GaussianFeH() : osink()
 	{
 		prov.insert("FeH");
+		prov.insert("FeH_comp{type=int; fmt=%3d;}");
 		req.insert("comp");
 		req.insert("XYZ");
 	}
 };
 extern "C" opipeline_stage *create_module_gaussianfeh() { return new os_GaussianFeH(); }	// Factory; called by opipeline_stage::create()
 
-DECLARE_KERNEL(os_GaussianFeH_kernel(otable_ks ks, bit_map applyToComponents, float mean, float sigma, gpu_rng_t rng,
+DECLARE_KERNEL(os_GaussianFeH_kernel(otable_ks ks, bit_map applyToComponents, os_GaussianFeH_data par, gpu_rng_t rng,
 		cint_t::gpu_t comp, cint_t::gpu_t hidden,
 		cfloat_t::gpu_t XYZ,
+		cint_t::gpu_t   FeH_comp,
 		cfloat_t::gpu_t FeH));
 
 size_t os_GaussianFeH::process(otable &in, size_t begin, size_t end, rng_t &rng)
@@ -65,21 +69,39 @@ size_t os_GaussianFeH::process(otable &in, size_t begin, size_t end, rng_t &rng)
 	cint_t   &hidden= in.col<int>("hidden");
 	cfloat_t &XYZ   = in.col<float>("XYZ");
 	cfloat_t &FeH   = in.col<float>("FeH");
+	cint_t   &FeH_comp = in.col<int>("FeH_comp");
 
-	CALL_KERNEL(os_GaussianFeH_kernel, otable_ks(begin, end), applyToComponents, mean, sigma, rng, comp, hidden, XYZ, FeH);
+	CALL_KERNEL(os_GaussianFeH_kernel, otable_ks(begin, end), applyToComponents, par, rng, comp, hidden, XYZ, FeH_comp, FeH);
 	return nextlink->process(in, begin, end, rng);
 }
 
 bool os_GaussianFeH::construct(const Config &cfg, otable &t, opipeline &pipe)
 {
 	read_component_map(applyToComponents, cfg);
 
-	mean  = cfg.get("mean");
-	sigma = cfg.get("sigma");
+	std::vector<float> mean  = cfg.get("mean");
+	std::vector<float> sigma = cfg.get("sigma");
+
+	ASSERT(mean.size() > 0);				// TODO: change to throw
+	ASSERT(mean.size()  < os_GaussianFeH_data::MAXCOMP);	// TODO: change to throw
+	ASSERT(sigma.size() == mean.size());			// TODO: change to throw
+	par.ncomp = mean.size();
+	FOR(0, par.ncomp) { par.mean[i] = mean[i]; }
+	FOR(0, par.ncomp) { par.sigma[i] = sigma[i]; }
+
+	// load the weights of each Gaussian, and compute the cumulative probability vector pcum
+	std::vector<double> wt = par.ncomp > 1 ? cfg.get("wt") : std::vector<double>(1, 1.);
+	ASSERT(wt.size() == par.ncomp);
+	std::partial_sum(wt.begin(), wt.end(), par.pcum);
+	ASSERT(par.pcum[par.ncomp-1] > 0);			// TODO: change to throw
+	FOR(0, par.ncomp) { par.pcum[i] /= par.pcum[par.ncomp-1]; }
 
 	// Output model parameters
-	MLOG(verb1) << "Metallicity: Gaussian for components " << applyToComponents << "   ## " << instanceName();
-	MLOG(verb2) << "           : (mu, sigma) = " << mean      << ", " << sigma;
+	MLOG(verb1) << "Metallicity: (multi)Gaussian for components " << applyToComponents << "   ## " << instanceName();
+	FOR(0, par.ncomp)
+	{
+		MLOG(verb2) << "           : (mu, sigma, pcum) = " << par.mean[i] << ", " << par.sigma[i] << ", " << par.pcum[i];
+	}
 
 	return true;
 }