Implementation of a Python functor for the target density function in…

… 'pyEXP' and 'exp'

expui/BiorthBasis.H

@@ -8,6 +8,7 @@
 #include <unsupported/Eigen/CXX11/Tensor> // For 3d rectangular grids
 #include <yaml-cpp/yaml.h>
 
+#include <DiskDensityFunc.H>
 #include <ParticleReader.H>
 #include <Coefficients.H>
 #include <BiorthBess.H>
@@ -592,16 +593,18 @@ namespace BasisClasses
     //! Basis construction parameters
     double aratio, hratio, dweight, rwidth, ashift, rfactor, rtrunc, ppow, Mfac, HERNA;
     bool Ignore, deproject;
+    std::string pyname;
 
     //! DiskType support
     //
-    enum DiskType { constant, gaussian, mn, exponential, doubleexpon, diskbulge };
+    enum DiskType { constant, gaussian, mn, exponential, doubleexpon, diskbulge, python };
     DiskType DTYPE;
     std::string mtype, dtype, dmodel;
 
     static const std::map<std::string, DiskType> dtlookup;
     double DiskDens(double R, double z, double phi);
     double dcond(double R, double z, double phi, int M);
+    std::shared_ptr<DiskDensityFunc> pyDens;
 
   protected:
 

expui/BiorthBasis.cc

@@ -60,6 +60,7 @@ namespace BasisClasses
     "self_consistent",
     "cachename",
     "modelname",
+    "pyname",
     "rnum"
   };
 
@@ -861,7 +862,8 @@ namespace BasisClasses
       {"mn",          DiskType::mn},
       {"exponential", DiskType::exponential},
       {"doubleexpon", DiskType::doubleexpon},
-      {"diskbulge", DiskType::diskbulge}
+      {"diskbulge",   DiskType::diskbulge},
+      {"python",      DiskType::python}
     };
 
   const std::set<std::string>
@@ -924,7 +926,8 @@ namespace BasisClasses
     "self_consistent",
     "playback",
     "coefCompute",
-    "coefMaster"
+    "coefMaster",
+    "pyname"
   };
 
   Cylindrical::Cylindrical(const YAML::Node& CONF) :
@@ -986,16 +989,19 @@ namespace BasisClasses
 
     case DiskType::diskbulge:
       {
-  double f = cosh(z/hcyl);
-  double rr = pow(pow(R, 2) + pow(z,2), 0.5);
-  double w1 = Mfac;
-  double w2 = (1-Mfac);
-  double as = HERNA;
-
-  ans = w1*exp(-R/acyl)/(4.0*M_PI*acyl*acyl*hcyl*f*f) + 
-            w2*pow(as, 4)/(2.0*M_PI*rr)*pow(rr+as,-3.0) ;
+	double f  = cosh(z/hcyl);
+	double rr = pow(pow(R, 2) + pow(z,2), 0.5);
+	double w1 = Mfac;
+	double w2 = 1.0 - Mfac;
+	double as = HERNA;
+	
+	ans = w1*exp(-R/acyl)/(4.0*M_PI*acyl*acyl*hcyl*f*f) + 
+	  w2*pow(as, 4)/(2.0*M_PI*rr)*pow(rr+as,-3.0) ;
       }
       break;
+    case DiskType::python:
+      ans = (*pyDens)(R, z, phi);
+      break;
     case DiskType::exponential:
     default:
       {
@@ -1078,7 +1084,6 @@ namespace BasisClasses
     // Radial scale length ratio for disk basis construction with doubleexpon
     aratio      = 1.0; 
 
-
     // Vertical scale height ratio for disk basis construction with doubleexpon
     hratio      = 1.0;              
 
@@ -1159,6 +1164,7 @@ namespace BasisClasses
       if (conf["mtype"     ])       mtype = conf["mtype"     ].as<std::string>();
       if (conf["dtype"     ])       dtype = conf["dtype"     ].as<std::string>();
       if (conf["vflag"     ])       vflag = conf["vflag"     ].as<int>();
+      if (conf["pyname"    ])      pyname = conf["pyname"    ].as<std::string>();
 
       // Deprecation warning
       if (conf["density"   ]) {
@@ -1298,6 +1304,12 @@ namespace BasisClasses
 	throw std::runtime_error("Cylindrical::initialize: invalid DiskType");
       }
 
+      // Check for and initialize the Python density type
+      //
+      if (DTYPE == DiskType::python) {
+	pyDens = std::make_shared<DiskDensityFunc>(pyname);
+      }
+
       // Use these user models to deproject for the EOF spherical basis
       //
       if (deproject) {

exputil/CMakeLists.txt

@@ -35,13 +35,14 @@ set(QPDISTF_SRC QPDistF.cc qld.c)
 set(SLEDGE_SRC sledge.f)
 set(PARTICLE_SRC Particle.cc ParticleReader.cc header.cc)
 set(CUDA_SRC cudaParticle.cu cudaSLGridMP2.cu)
+set(PYWRAP_SRC DiskDensityFunc.cc)
 
 set(exputil_SOURCES ${ODE_SRC} ${ROOT_SRC} ${QUAD_SRC}
   ${RANDOM_SRC} ${UTIL_SRC} ${SPECFUNC_SRC}
   ${PHASE_SRC}  ${SYMP_SRC} ${INTERP_SRC} ${MASSMODEL_SRC}
   ${ORBIT_SRC} ${BIORTH_SRC} ${POLY_SRC} ${GAUSS_SRC}
   ${QPDISTF_SRC} ${BESSEL_SRC} ${OPTIMIZATION_SRC}
-  ${SLEDGE_SRC} ${PARTICLE_SRC} ${CUDA_SRC})
+  ${SLEDGE_SRC} ${PARTICLE_SRC} ${CUDA_SRC} ${PYWRAP_SRC})
 
 set(common_INCLUDE_DIRS $<INSTALL_INTERFACE:include>
   $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/include/>
@@ -53,7 +54,7 @@ set(common_INCLUDE_DIRS $<INSTALL_INTERFACE:include>
 
 set(common_LINKLIB ${DEP_LIB} OpenMP::OpenMP_CXX MPI::MPI_CXX yaml-cpp
   ${VTK_LIBRARIES} ${HDF5_CXX_LIBRARIES} ${HDF5_LIBRARIES}
-  ${HDF5_HL_LIBRARIES} ${FFTW_DOUBLE_LIB})
+  ${HDF5_HL_LIBRARIES} ${FFTW_DOUBLE_LIB} pybind11::embed)
 
 if(ENABLE_CUDA)
   list(APPEND common_LINKLIB CUDA::toolkit CUDA::cudart)

include/DiskDensityFunc.H

@@ -0,0 +1,33 @@
+#ifndef _DiskDensityFunc_H_
+#define _DiskDensityFunc_H_
+
+#include <functional>
+
+#include <pybind11/pybind11.h>
+#include <pybind11/embed.h>
+
+// Convenience
+namespace py = pybind11;
+
+//! Python disk-density function wrapper
+class DiskDensityFunc
+{
+private:
+
+  //! The disk-density function
+  py::function disk_density;
+
+public:
+
+  //! Constructor
+  DiskDensityFunc(const std::string& modulename);
+
+  //! Destructor
+  ~DiskDensityFunc();
+
+  //! Evaluate the density
+  double operator() (double R, double z, double phi);
+
+};
+
+#endif

include/DiskModels.H

@@ -2,7 +2,7 @@
 #define _DISK_MODELS_H
 
 #include <EmpCylSL.H>
-
+#include <DiskDensityFunc.H>
 
 //! A Ferrers Ellipsoid + Evacuated Exponential Disc (semi-realistic
 //! bar+disk model)
@@ -255,4 +255,26 @@ public:
 
 };
 
-#endif
+
+//! A user-defined Python model
+class PyModel : public EmpCylSL::AxiDisk
+{
+private:
+
+  std::shared_ptr<DiskDensityFunc> pyDens;
+
+public:
+
+  PyModel(std::string& pyname)
+  {
+    pyDens = std::make_shared<DiskDensityFunc>(pyname);
+  }
+
+  double operator()(double R, double z, double phi=0.)
+  {
+    return (*pyDens)(R, z, phi);
+  }
+
+};
+
+#endif

src/Cylinder.H

@@ -103,6 +103,8 @@ class MixtureBasis;
 
     @param playback file reads a coefficient file and uses it to compute the basis function output for resimiulation
 
+    @param python is the file name of Python module which supplies the 'disk_density' function for conditioning the cylindrical basis.  A non-null string triggers the use of the Python interpreter to evaluate the target density function.
+
 */
 class Cylinder : public Basis
 {
@@ -130,7 +132,7 @@ private:
   int ncylnx, ncylny, ncylr;
   double hcyl, hexp, snr, rem;
   int nmax, ncylodd, ncylrecomp, npca, npca0, nvtk, cmapR, cmapZ;
-  std::string cachename;
+  std::string cachename, pyname;
   bool self_consistent, logarithmic, pcavar, pcainit, pcavtk, pcadiag, pcaeof;
   bool try_cache, firstime, dump_basis, compute, firstime_coef;
 

src/Cylinder.cc

@@ -10,52 +10,17 @@
 #include <CylEXP.H>
 #include <Cylinder.H>
 #include <MixtureBasis.H>
+#include <DiskDensityFunc.H>
 #include <Timer.H>
 #include <exputils.H>
 #include <NVTX.H>
 
-Timer timer_debug;
-
-double EXPSCALE=1.0, HSCALE=1.0, ASHIFT=0.25;
-
-
 //@{
 //! These are for testing exclusively (should be set false for production)
 static bool cudaAccumOverride = false;
 static bool cudaAccelOverride = false;
 //@}
 
-double DiskDens(double R, double z, double phi)
-{
-  double f = cosh(z/HSCALE);
-  return exp(-R/EXPSCALE)/(4.0*M_PI*EXPSCALE*EXPSCALE*HSCALE*f*f);
-}
-
-
-double dcond(double R, double z, double phi, int M)
-{
-  //
-  // No shift for M==0
-  //
-  if (M==0) return DiskDens(R, z, phi);
-
-  //
-  // Fold into [-PI/M, PI/M] for M>=1
-  //
-  double dmult = M_PI/M, phiS;
-  if (phi>M_PI)
-    phiS = phi + dmult*(int)((2.0*M_PI - phi)/dmult);
-  else
-    phiS = phi - dmult*(int)(phi/dmult);
-
-  //
-  // Apply a shift along the x-axis
-  //
-  double x = R*cos(phiS) - ASHIFT*EXPSCALE;
-  double y = R*sin(phiS);
-  return DiskDens(sqrt(x*x + y*y), z, atan2(y, x));
-}
-
 
 const std::set<std::string>
 Cylinder::valid_keys = {
@@ -106,7 +71,8 @@ Cylinder::valid_keys = {
   "self_consistent",
   "playback",
   "coefCompute",
-  "coefMaster"
+  "coefMaster",
+  "pyname"
 };
 
 Cylinder::Cylinder(Component* c0, const YAML::Node& conf, MixtureBasis *m) :
@@ -237,9 +203,6 @@ Cylinder::Cylinder(Component* c0, const YAML::Node& conf, MixtureBasis *m) :
 
     // Set parameters for external dcond function
     //
-    EXPSCALE = acyl;
-    HSCALE   = hcyl;
-    ASHIFT   = ashift;
     eof      = 0;
 
     bool cache_ok = false;
@@ -299,6 +262,40 @@ Cylinder::Cylinder(Component* c0, const YAML::Node& conf, MixtureBasis *m) :
 		  << "this step will take many minutes."
 		  << std::endl;
 
+      std::shared_ptr<DiskDensityFunc> dfunc;
+      if (pyname.size()) dfunc = std::make_shared<DiskDensityFunc>(pyname);
+
+      auto DiskDens = [&](double R, double z, double phi)
+      {
+	if (dfunc) return (*dfunc)(R, z, phi);
+	double f = cosh(z/hcyl);
+	return exp(-R/acyl)/(4.0*M_PI*acyl*acyl*hcyl*f*f);
+      };
+
+      auto dcond = [&](double R, double z, double phi, int M)
+      {
+	//
+	// No shift for M==0
+	//
+	if (M==0) return DiskDens(R, z, phi);
+
+	//
+	// Fold into [-PI/M, PI/M] for M>=1
+	//
+	double dmult = M_PI/M, phiS;
+	if (phi>M_PI)
+	  phiS = phi + dmult*(int)((2.0*M_PI - phi)/dmult);
+	else
+	  phiS = phi - dmult*(int)(phi/dmult);
+
+	//
+	// Apply a shift along the x-axis
+	//
+	double x = R*cos(phiS) - ashift*acyl;
+	double y = R*sin(phiS);
+	return DiskDens(sqrt(x*x + y*y), z, atan2(y, x));
+      };
+
       ortho->generate_eof(rnum, pnum, tnum, dcond);
     }
 

utils/ICs/cylcache.cc

@@ -165,18 +165,20 @@ void set_fpu_gdb_handler(void)
 
 // Global variables
 //
-enum DiskType { constant, gaussian, mn, exponential, doubleexpon, diskbulge };
+enum DiskType { constant, gaussian, mn, exponential, doubleexpon, diskbulge, python };
 
 std::map<std::string, DiskType> dtlookup =
   { {"constant",    DiskType::constant},
     {"gaussian",    DiskType::gaussian},
     {"mn",          DiskType::mn},
     {"exponential", DiskType::exponential},
     {"doubleexpon", DiskType::doubleexpon},
-    {"diskbulge", DiskType::diskbulge}
+    {"diskbulge",   DiskType::diskbulge},
+    {"python",      DiskType::python}
   };
 
 DiskType     DTYPE;
+std::string  pyname;
 double       ASCALE;
 double       ASHIFT;
 double       HSCALE;
@@ -190,6 +192,8 @@ double       DWEIGHT = 1.0;
 
 #include <Particle.H>
 
+static std::shared_ptr<DiskDensityFunc> pydens;
+
 double DiskDens(double R, double z, double phi)
 {
   double ans = 0.0;
@@ -245,6 +249,11 @@ double DiskDens(double R, double z, double phi)
             w2*pow(as, 4)/(2.0*M_PI*rr)*pow(rr+as,-3.0) ; 
     }
     break;
+  case DiskType::python:
+    {
+      if (not pydens) pydens = std::make_shared<DiskDensityFunc>(pyname);
+      ans = (*pydens)(R, z, phi);
+    }
   case DiskType::exponential:
   default:
     {
@@ -407,6 +416,8 @@ main(int ac, char **av)
      cxxopts::value<string>(dmodel)->default_value("0.1"))
     ("Mfac", "Fraction of mass in the disk, remaining mass will be in the bulge (i.e. Mfac = .75 would have 75% of the mass be disk, 25% bulge). diskbulge model only",
      cxxopts::value<string>(dmodel)->default_value("0.1"))
+    ("pyname", "The name of the python module supplying the disk_density",
+     cxxopts::value<std::string>(pyname))
     ;
 
   cxxopts::ParseResult vm;