From 5a0fd09fb785b010e8ca2b47c3d4f30bcf6ca041 Mon Sep 17 00:00:00 2001
From: Eduard Valeyev <eduard@valeyev.net>
Date: Mon, 11 Sep 2023 06:07:38 -0400
Subject: [PATCH] introduced Function::print_tree_json

---
 src/madness/mra/funcimpl.h | 28 ++++++++++++++++++++
 src/madness/mra/mra.h      |  7 +++++
 src/madness/mra/mraimpl.h  | 54 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+)

diff --git a/src/madness/mra/funcimpl.h b/src/madness/mra/funcimpl.h
index 2befa0da804..2c89e977652 100644
--- a/src/madness/mra/funcimpl.h
+++ b/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 {
diff --git a/src/madness/mra/mra.h b/src/madness/mra/mra.h
index 0b7f515bcdb..4c3b2fb3875 100644
--- a/src/madness/mra/mra.h
+++ b/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);
diff --git a/src/madness/mra/mraimpl.h b/src/madness/mra/mraimpl.h
index 283d060803b..c3ef462a863 100644
--- a/src/madness/mra/mraimpl.h
+++ b/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();