Evaleev/feature/wolfram orbital plotting

src/madness/mra/funcimpl.h

@@ -450,6 +450,26 @@ namespace madness {
             ar & coeff() & _has_children & _norm_tree & dnorm & snorm;
         }
 
+        /// like operator<<(ostream&, const FunctionNode<T,NDIM>&) but
+        /// produces a sequence JSON-formatted key-value pairs
+        /// @warning enclose the output in curly braces to make
+        /// a valid JSON object
+        void print_json(std::ostream& s) const {
+            s << "\"has_coeff\":" << this->has_coeff()
+              << ",\"has_children\":" << this->has_children() << ",\"norm\":";
+            double norm = this->has_coeff() ? this->coeff().normf() : 0.0;
+            if (norm < 1e-12)
+                norm = 0.0;
+            double nt = this->get_norm_tree();
+            if (nt == 1e300)
+                nt = 0.0;
+            s << norm << ",\"norm_tree\":" << nt << ",\"snorm\":"
+              << this->get_snorm() << ",\"dnorm\":" << this->get_dnorm()
+              << ",\"rank\":" << this->coeff().rank();
+            if (this->coeff().is_assigned())
+                s << ",\"dim\":" << this->coeff().dim(0);
+        }
+
     };
 
     template <typename T, std::size_t NDIM>
@@ -1311,6 +1331,14 @@ namespace madness {
         /// Functor for the do_print_tree method (using GraphViz)
         void do_print_tree_graphviz(const keyT& key, std::ostream& os, Level maxlevel) const;
 
+        /// Same as print_tree() but in JSON format
+        /// @param[out] os the ostream to where the output is sent
+        /// @param[in] maxlevel the maximum level of the tree for printing
+        void print_tree_json(std::ostream& os = std::cout, Level maxlevel = 10000) const;
+
+        /// Functor for the do_print_tree_json method
+        void do_print_tree_json(const keyT& key, std::multimap<Level, std::tuple<tranT, std::string>>& data, Level maxlevel) const;
+
         /// convert a number [0,limit] to a hue color code [blue,red],
         /// or, if log is set, a number [1.e-10,limit]
         struct do_convert_to_color {

src/madness/mra/funcplot.h

@@ -952,30 +952,67 @@ void plot_plane(World& world, const std::vector<Function<double,NDIM> >& vfuncti
         fprintf(file,"%d %12.8f %12.8f %12.8f \n",int(molecular_info.size()),
                 cell(0,0),cell(1,0),cell(2,0));
 
+        // grid spacing for each dimension such that the cell edges are plotted
+        const auto xdelta = xlen/(npt[0]-1);
+        const auto ydelta = ylen/(npt[1]-1);
+        const auto zdelta = zlen/(npt[2]-1);
+
         // print the cell constants
-        fprintf(file,"%d %12.6f %12.6f %12.6f\n",npt[0],xlen/npt[0],0.0,0.0);
-        fprintf(file,"%d %12.6f %12.6f %12.6f\n",npt[1],0.0,ylen/npt[1],0.0);
-        fprintf(file,"%d %12.6f %12.6f %12.6f\n",npt[2],0.0,0.0,zlen/npt[2]);
+        fprintf(file,"%d %12.6f %12.6f %12.6f\n",npt[0],xdelta,0.0,0.0);
+        fprintf(file,"%d %12.6f %12.6f %12.6f\n",npt[1],0.0,ydelta,0.0);
+        fprintf(file,"%d %12.6f %12.6f %12.6f\n",npt[2],0.0,0.0,zdelta);
 
         // print the molecule
         for (const std::string& s : molecular_info) fprintf(file,"%s",s.c_str());
 
 
-        Tensor<double>grid(npt[0],npt[1],npt[2]);
-        for (int i=0;i<npt[0];++i) {
-            for (int j=0;j<npt[1];++j) {
-                for (int k=0;k<npt[2];++k) {
-                    double x=cell(0,0)+origin[0]+xlen/npt[0]*i;
-                    double y=cell(1,0)+origin[1]+ylen/npt[1]*j;
-                    double z=cell(2,0)+origin[2]+zlen/npt[2]*k;
-                    fprintf(file,"%12.8f",f(x,y,z));
-                }
-            }
-            fprintf(file,"\n");
-        }
-        fclose(file);
+         Tensor<double> grid(npt[0], npt[1], npt[2]);
+         long count_per_line = 0;
+         for (int i = 0; i < npt[0]; ++i) {
+             for (int j = 0; j < npt[1]; ++j) {
+                 for (int k = 0; k < npt[2]; ++k) {
+                     double x = cell(0, 0) + origin[0] + xdelta * i;
+                     double y = cell(1, 0) + origin[1] + ydelta * j;
+                     double z = cell(2, 0) + origin[2] + zdelta * k;
+                     // the original format has up to 6 entries per line: https://paulbourke.net/dataformats/cube/
+                     // stick with this, even though many codes can read an arbitrary number of entries per line
+                     if (count_per_line < 6) {
+                        fprintf(file, "%12.5e ", f(x, y, z));
+                        count_per_line++;
+                     } else {
+                        fprintf(file, "%12.5e\n", f(x, y, z));
+                        count_per_line = 0;
+                     }
+                 }
+             }
+         }
+         fprintf(file, "\n");
+         fclose(file);
+     }
 
-    }
+     template<size_t NDIM>
+     typename std::enable_if<NDIM==3,void>::type
+     print_tree_jsonfile(World& world, const Function<double,NDIM>& f, std::string filename) {
+
+         if (world.size() > 1)
+             return;
+
+         Tensor<double> cell = copy(FunctionDefaults<3>::get_cell());
+         std::ofstream os(filename.c_str());
+
+         os << "{";
+         os << "\"cell\":[";
+         for (int xyz = 0; xyz != 3; ++xyz) {
+             os << "[" << cell(xyz, 0) << "," << cell(xyz, 1) << "]";
+             if (xyz != 2)
+                 os << ",";
+         }
+         os << "],";
+
+         os << "\"tree\":{";
+         f.print_tree_json(os);
+         os << "}}";
+     }
 
     /// convenience to get plot_plane and plot_cubefile
     template<size_t NDIM>

src/madness/mra/mra.h

@@ -859,6 +859,13 @@ namespace madness {
             if (impl) impl->print_tree(os);
         }
 
+        /// same as print_tree() but produces JSON-formatted string
+        /// @warning enclose the result in braces to make it a valid JSON object
+        void print_tree_json(std::ostream& os = std::cout) const {
+            PROFILE_MEMBER_FUNC(Function);
+            if (impl) impl->print_tree_json(os);
+        }
+
         /// Process 0 prints a graphviz-formatted output of all nodes in the tree (collective)
         void print_tree_graphviz(std::ostream& os = std::cout) const {
             PROFILE_MEMBER_FUNC(Function);

src/madness/mra/mraimpl.h

@@ -2706,8 +2706,62 @@ namespace madness {
         }
     }
 
+    template <typename T, std::size_t NDIM>
+    void FunctionImpl<T,NDIM>::print_tree_json(std::ostream& os, Level maxlevel) const {
+        std::multimap<Level, std::tuple<tranT, std::string>> data;
+        if (world.rank() == 0) do_print_tree_json(cdata.key0, data, maxlevel);
+        world.gop.fence();
+        if (world.rank() == 0) {
+            for (Level level = 0; level != maxlevel; ++level) {
+                if (data.count(level) == 0)
+                    break;
+                else {
+                    if (level > 0)
+                        os << ",";
+                    os << "\"" << level << "\":{";
+                    os << "\"level\": " << level << ",";
+                    os << "\"nodes\":{";
+                    auto range = data.equal_range(level);
+                    for (auto it = range.first; it != range.second; ++it) {
+                        os << "\"" << std::get<0>(it->second) << "\":"
+                           << std::get<1>(it->second);
+                        if (std::next(it) != range.second)
+                            os << ",";
+                    }
+                    os << "}}";
+                }
+            }
+            os.flush();
+        }
+        world.gop.fence();
+    }
+
+
+    template <typename T, std::size_t NDIM>
+    void FunctionImpl<T,NDIM>::do_print_tree_json(const keyT& key, std::multimap<Level, std::tuple<tranT, std::string>>& data, Level maxlevel) const {
+        typename dcT::const_iterator it = coeffs.find(key).get();
+        if (it == coeffs.end()) {
+            MADNESS_EXCEPTION("FunctionImpl: do_print_tree_json: null node pointer",0);
+        }
+        else {
+            const nodeT& node = it->second;
+            std::ostringstream oss;
+            oss << "{";
+            node.print_json(oss);
+            oss << ",\"owner\": " << coeffs.owner(key) << "}";
+            auto node_json_str = oss.str();
+            data.insert(std::make_pair(key.level(), std::make_tuple(key.translation(), node_json_str)));
+            if (key.level() < maxlevel  &&  node.has_children()) {
+                for (KeyChildIterator<NDIM> kit(key); kit; ++kit) {
+                    do_print_tree_json(kit.key(),data, maxlevel);
+                }
+            }
+        }
+    }
+
     template <typename T, std::size_t NDIM>
     void FunctionImpl<T,NDIM>::print_tree_graphviz(std::ostream& os, Level maxlevel) const {
+        // aggregate data by level, thus collect data first, then dump
         if (world.rank() == 0) do_print_tree_graphviz(cdata.key0, os, maxlevel);
         world.gop.fence();
         if (world.rank() == 0) os.flush();

src/madness/mra/tools/MRAMeshOrbitalPlot3D.wl

@@ -0,0 +1,118 @@
+(* ::Package:: *)
+
+BeginPackage["MRAMeshOrbitalPlot3D`"]
+
+
+(* ::Text:: *)
+(*{*)
+(* {Description: MRAMeshOrbitalPlot3D[fp,opts] returns a plot of the orbital stored via madness::plot_cubefile and madness::print_json_treefile in files fp . cube and fp . tree . json, respectively . MRAMeshOrbitalPlot3D accepts all options recognized by Graphics3D and ListContourPlot3D functions, as well as the following additional options:, \[SpanFromLeft], \[SpanFromLeft]},*)
+(* {Option, Default, Description},*)
+(* {Zoom, 1, can be used to produce a plot in a zoomed-in section of the simulation cell . This does not need to match the zoom value given to plot_cubefile (that only affects the resolution/extent of the Gaussian Cube mesh)},*)
+(* {MRAMeshCuboidDirectives, {EdgeForm[Thick]}, Specifies how the Cuboid objects comprising the MRA mesh are drawn . All Graphics3D directives that are applicable to Cuboid (except Opacity) can be specified . },*)
+(* {MaxLevel, Infinity, Controls the highest refinement level of displayed mesh elements .},*)
+(* {MinLevel, 0, Controls the lowest refinement level of displayed mesh elements .}*)
+(*}*)
+
+
+MRAMeshOrbitalPlot3D::usage="MRAMeshOrbitalPlot3D[fp,opts] returns a plot of the orbital stored via madness::plot_cubefile and madness::print_json_treefile in files fp.cube and fp.tree.json, respectively. MRAMeshOrbitalPlot3D accepts all options recognized by Graphics3D and ListContourPlot3D functions, as well as the following additional options: Zoom, MRAMeshCuboidDirectives, MinLevel, and MaxLevel."
+
+
+Begin["`Private`"]
+
+
+ReadTree[fileName_] :=
+    Module[{jsonData, cellData, treeData, boxCoords, nodesData},
+        jsonData = Import[fileName];
+        cellData = "cell" /. jsonData;
+        treeData = "tree" /. jsonData;
+        boxCoords = {};
+        Do[
+            nodesData = "nodes" /. (ToString[n] /. treeData);
+            Do[AppendTo[boxCoords, Prepend[Interpreter["Integer"] /@ 
+                StringSplit[Part[nodesData[[i]], 1], {"[", "]", ","}], n]], {i, Length[
+                nodesData]}];
+            ,
+            {n, 0, Length[treeData] - 1}
+        ];
+        Return[{cellData, boxCoords}]
+    ];
+
+BoxToXYZCoords[cell_, box_] :=
+    Module[{S},
+        S = Table[cell[[xyz, 2]] - cell[[xyz, 1]], {xyz, 3}];
+        Return[{{cell[[1, 1]] + S[[1]] box[[2]] / (2 ^ box[[1]]), cell
+            [[2, 1]] + S[[2]] box[[3]] / (2 ^ box[[1]]), cell[[3, 1]] + S[[3]] box
+            [[4]] / (2 ^ box[[1]])}, {cell[[1, 1]] + S[[1]] (box[[2]] + 1) / (2 ^
+             box[[1]]), cell[[2, 1]] + S[[2]] (box[[3]] + 1) / (2 ^ box[[1]]), cell
+            [[3, 1]] + S[[3]] (box[[4]] + 1) / (2 ^ box[[1]])}}]
+    ];
+
+BoxToGraphics[cell_, box_, nmax_, omax_(* opacity value of the smallest
+     boxes *), shadeexp_(* opacity of box at level n-1 is this times smaller
+     than than of box at level n *),cuboidDirectives_List] :=
+    Module[{},
+        Return[Join[(*N.B. MUST BE FIRST to apply to Cuboid*)cuboidDirectives,{Opacity[omax / shadeexp ^ (nmax - box[[1]])], Cuboid 
+            @@ BoxToXYZCoords[cell, box],EdgeForm[Thick]}]]
+    ];
+
+(*
+MaxLevel,MinLevel: show boxes with resolution level [nmin,nmax]
+Zoom: limit PlotRange to cell/Zoom
+*)
+
+Protect[MaxLevel,MinLevel,Zoom,MRAMeshCuboidDirectives];
+Options[MRAMeshOrbitalPlot3D] = {MaxLevel -> Infinity, MinLevel -> 0, Zoom
+     -> 1,MRAMeshCuboidDirectives->{EdgeForm[Thick]}};
+
+(* plots orbital and its mesh superimposed *)
+
+MRAMeshOrbitalPlot3D[filePrefix_, opt : OptionsPattern[{MRAMeshOrbitalPlot3D,
+     Graphics3D, ListContourPlot3D, Show}]] :=
+    Module[
+        {simulationCell, boxTreeCoords, bohr2angstrom, selectedBoxes,
+             boxGraphics, meshPlot, orbitalPlot, nmax, nmin, zoom, omax, shadeexp,
+             actualNMax,cuboidDirectives, plotRange}
+        ,
+        (* process options *)
+        nmax = OptionValue[MaxLevel];
+        nmin = OptionValue[MinLevel];
+        zoom = OptionValue[Zoom];
+        cuboidDirectives=OptionValue[MRAMeshCuboidDirectives];
+
+        (* these are hardwired since they are not needed for most users
+             *)
+        omax = 0;(* opacity value of the smallest boxes *)
+        shadeexp = 1.9;(* opacity of box at level n-1 is this times smaller
+             than than of box at level n *)
+
+        {simulationCell, boxTreeCoords} = ReadTree[filePrefix <> ".tree.json"
+            ]; (* the deduced hardwired value used by Wolfram when importing Gaussian
+             Cube file *) bohr2angstrom = 0.529177249;
+        simulationCell *= bohr2angstrom;
+        (* override default PlotRange *)
+        plotRange =
+            If[OptionValue[PlotRange] === All,
+                simulationCell / zoom
+                ,
+                OptionValue[PlotRange]
+            ];
+        selectedBoxes = Select[boxTreeCoords, #[[1]] <= nmax && #[[1]]
+             >= nmin&];
+        actualNMax = MaximalBy[selectedBoxes, #[[1]]&][[1, 1]];
+        boxGraphics = Map[BoxToGraphics[simulationCell, #, actualNMax,
+             omax, shadeexp, cuboidDirectives]&, selectedBoxes];
+        meshPlot = Graphics3D[boxGraphics, Boxed -> False, Evaluate @
+             FilterRules[{opt}, Options[Graphics3D]]];
+        orbitalPlot = Import[filePrefix <> ".cube", "Graphics3D", Boxed
+             -> False, Evaluate @ FilterRules[{opt}, Options[ListContourPlot3D]]]
+            ;
+        Return[Show[{orbitalPlot, meshPlot}, Evaluate @ FilterRules[{
+            opt, PlotRange -> plotRange}, Options[Graphics3D]]]];
+    ];
+
+
+
+End[]
+
+
+EndPackage[]

src/madness/mra/tools/README.md

@@ -0,0 +1,37 @@
+# Contents
+
+This directory contains various utilities:
+
+- autocorr.mw ... a Maple worksheet to generate the autocorrelation coefficients
+- dump2.py with dependencies ... a Python program to generate the two-scale coefficients. Run with `python dump2.py`
+- quadrature.py can be easily modified to generate the full-precision Gauss-Legendre coeffs
+- MRAMeshOrbitalPlot3D.wl ... a Wolfram language ("Mathematica", for the old school) package for plotting orbitals stored in the Gaussian Cube format (see `madness::plot_cubefile`) along with the mesh stored a JSON format (see `madness::print_tree_jsonfile`).
+
+## `MRAMeshOrbitalPlot3D.wl`
+
+- if you have [WolframScript](https://www.wolfram.com/wolframscript/) installed (it is bundled with Mathematica by default on Windows/Linux; on a Mac need to [download/install manually](https://www.wolfram.com/wolframscript/) to plot from command-line or shell script. See a quick example in `h2-no1.wsl`; to try out execute `./h2-no1.wsl > h2-no1.pdf` 
+- more elaborate plotting is best done from a Mathematica notebook. Here's a brief example:
+```Wolfram
+<< "/path/to/madness/source/dir/src/madness/mra/tools/MRAMeshOrbitalPlot3D.wl"
+
+(* this assumes that data.cube (produced by madness::plot_cubefile) and data.tree.json (produced by madness::print_tree_jsonfile) exist in /path/to/data
+MRAMeshOrbitalPlot3D["/path/to/data"]
+```
+You can even pull the module and data remotely, e.g.:
+```Wolfram
+<< "https://raw.githubusercontent.com/m-a-d-n-e-s-s/madness/evaleev/feature/wolfram-orbital-plotting/src/madness/mra/tools/MRAMeshOrbitalPlot3D.wl"
+
+MRAMeshOrbitalPlot3D["https://raw.githubusercontent.com/m-a-d-n-e-s-s/madness/evaleev/feature/wolfram-orbital-plotting/src/madness/mra/tools/h2-no1"]
+```
+
+### 
+The only function `MRAMeshOrbitalPlot3D.wl` provides is `MRAMeshOrbitalPlot3D`.  `MRAMeshOrbitalPlot3D[fp,opts]` returns a plot of the orbital stored via `madness::plot_cubefile` and `madness::print_json_treefile` in files `fp.cube` and `fp.tree.json`, respectively. `MRAMeshOrbitalPlot3D` accepts all options recognized by [`Graphics3D`](https://reference.wolfram.com/language/ref/Graphics3D.html]) and [`ListContourPlot3D`](https://reference.wolfram.com/language/ref/ListContourPlot3D.html) functions, as well as the following additional options:
+
+| Option                      | Default             | Description                                                                                                                                                                                                                                                                                                                                             |
+|-----------------------------|---------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `Zoom`                      | `1`                 | Can be used to produce a plot in a zoomed-in section of the simulation cell. This does not need to match the zoom value given to plot_cubefile (that only affects the resolution/extent of the Gaussian Cube mesh)                                                                                                                                      |
+| `MRAMeshCuboidDirectives`   | `{EdgeForm[Thick]}` | Specifies how the Cuboid objects comprising the MRA mesh are drawn. All [`Graphics3D`](https://reference.wolfram.com/language/ref/Graphics3D.html) directives that are applicable to [`Cuboid`](https://reference.wolfram.com/language/ref/Cuboid.html) (except [`Opacity`](https://reference.wolfram.com/language/ref/Opacity.html)) can be specified. |
+| `MaxLevel`                  | `Infinity`          | Controls the highest refinement level of displayed mesh elements.                                                                                                                                                                                                                                                                                       |
+| `MinLevel`                  | `0`                 | Controls the lowest refinement level of displayed mesh elements.                                                                                                                                                                                                                                                                                        |
+
+