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&) 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 @@ -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>& 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 + void FunctionImpl::print_tree_json(std::ostream& os, Level maxlevel) const { + std::multimap> 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 + void FunctionImpl::do_print_tree_json(const keyT& key, std::multimap>& 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 kit(key); kit; ++kit) { + do_print_tree_json(kit.key(),data, maxlevel); + } + } + } + } + template void FunctionImpl::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();