Implement fetch_sum and fetch_nnz

paulsengroup · Jul 30, 2023 · 80c57f6 · 80c57f6
1 parent 92b21ce
commit 80c57f6
Show file tree

Hide file tree

Showing 5 changed files with 93 additions and 1 deletion.
diff --git a/src/common.hpp b/src/common.hpp
@@ -407,4 +407,44 @@ inline py::object file_fetch_dense(const File &f, std::string_view range1, std::
                                   num_rows, num_cols, bin1.id(), bin2.id());
 }
 
+template <typename File>
+inline py::object file_fetch_sum(const File &f, std::string_view range1, std::string_view range2,
+                                 std::string_view normalization, std::string_view count_type,
+                                 std::string_view query_type) {
+  if (normalization != "NONE") {
+    count_type = "float";
+  }
+
+  const auto qt =
+      query_type == "UCSC" ? hictk::GenomicInterval::Type::UCSC : hictk::GenomicInterval::Type::BED;
+
+  auto sel = range2.empty() || range1 == range2
+                 ? f.fetch(range1, hictk::balancing::Method(normalization), qt)
+                 : f.fetch(range1, range2, hictk::balancing::Method(normalization), qt);
+
+  if (count_type == "int") {
+    return py::cast(std::accumulate(
+        sel.template begin<std::int32_t>(), sel.template end<std::int32_t>(), std::int64_t(0),
+        [](const auto accumulator, const hictk::ThinPixel<std::int32_t> &p) {
+          return accumulator + p.count;
+        }));
+  }
+  return py::cast(std::accumulate(sel.template begin<double>(), sel.template end<double>(), 0.0,
+                                  [](const auto accumulator, const hictk::ThinPixel<double> &p) {
+                                    return accumulator + p.count;
+                                  }));
+}
+
+template <typename File>
+inline std::int64_t file_fetch_nnz(const File &f, std::string_view range1, std::string_view range2,
+                                   std::string_view query_type) {
+  const auto qt =
+      query_type == "UCSC" ? hictk::GenomicInterval::Type::UCSC : hictk::GenomicInterval::Type::BED;
+
+  auto sel = range2.empty() || range1 == range2
+                 ? f.fetch(range1, hictk::balancing::Method("NONE"), qt)
+                 : f.fetch(range1, range2, hictk::balancing::Method("NONE"), qt);
+  return std::distance(sel.template begin<std::int32_t>(), sel.template end<std::int32_t>());
+}
+
 }  // namespace hictkpy
diff --git a/src/hictkpy.cpp b/src/hictkpy.cpp
@@ -51,6 +51,11 @@ static pybind11::module_ declare_cooler_submodule(pybind11::module_ &m) {
   cooler_file.def("fetch_dense", &cooler::fetch_dense, py::arg("range1") = "",
                   py::arg("range2") = "", py::arg("normalization") = "NONE",
                   py::arg("count_type") = "int", py::arg("query_type") = "UCSC");
+  cooler_file.def("fetch_sum", &cooler::fetch_sum, py::arg("range1") = "", py::arg("range2") = "",
+                  py::arg("normalization") = "NONE", py::arg("count_type") = "int",
+                  py::arg("query_type") = "UCSC");
+  cooler_file.def("fetch_nnz", &cooler::fetch_nnz, py::arg("range1") = "", py::arg("range2") = "",
+                  py::arg("query_type") = "UCSC");
 
   return cooler;
 }
@@ -86,7 +91,11 @@ static pybind11::module_ declare_hic_submodule(pybind11::module_ &m) {
   hic_file.def("fetch_dense", &hic::fetch_dense, py::arg("range1") = "", py::arg("range2") = "",
                py::arg("normalization") = "NONE", py::arg("count_type") = "int",
                py::arg("query_type") = "UCSC");
-
+  hic_file.def("fetch_sum", &hic::fetch_sum, py::arg("range1") = "", py::arg("range2") = "",
+               py::arg("normalization") = "NONE", py::arg("count_type") = "int",
+               py::arg("query_type") = "UCSC");
+  hic_file.def("fetch_nnz", &hic::fetch_nnz, py::arg("range1") = "", py::arg("range2") = "",
+               py::arg("query_type") = "UCSC");
   return hic;
 }
 
@@ -117,6 +126,11 @@ static void declare_file_class(pybind11::module_ &m) {
   file.def("fetch_dense", &file::fetch_dense, py::arg("range1") = "", py::arg("range2") = "",
            py::arg("normalization") = "NONE", py::arg("count_type") = "int",
            py::arg("query_type") = "UCSC");
+  file.def("fetch_sum", &file::fetch_sum, py::arg("range1") = "", py::arg("range2") = "",
+           py::arg("normalization") = "NONE", py::arg("count_type") = "int",
+           py::arg("query_type") = "UCSC");
+  file.def("fetch_nnz", &file::fetch_nnz, py::arg("range1") = "", py::arg("range2") = "",
+           py::arg("query_type") = "UCSC");
 }
 
 namespace py = pybind11;

diff --git a/src/hictkpy_cooler.hpp b/src/hictkpy_cooler.hpp
@@ -120,4 +120,15 @@ inline py::object fetch_dense(const hictk::cooler::File &f, std::string_view ran
                               std::string_view count_type, std::string_view query_type) {
   return file_fetch_dense(f, range1, range2, normalization, count_type, query_type);
 }
+
+inline py::object fetch_sum(const hictk::cooler::File &f, std::string_view range1,
+                            std::string_view range2, std::string_view normalization,
+                            std::string_view count_type, std::string_view query_type) {
+  return file_fetch_sum(f, range1, range2, normalization, count_type, query_type);
+}
+
+inline std::int64_t fetch_nnz(const hictk::cooler::File &f, std::string_view range1,
+                              std::string_view range2, std::string_view query_type) {
+  return file_fetch_nnz(f, range1, range2, query_type);
+}
 }  // namespace hictkpy::cooler
diff --git a/src/hictkpy_file.hpp b/src/hictkpy_file.hpp
@@ -49,4 +49,20 @@ inline py::object fetch_dense(const hictk::File &f, std::string_view range1,
       f.get());
 }
 
+inline py::object fetch_sum(const hictk::File &f, std::string_view range1, std::string_view range2,
+                            std::string_view normalization, std::string_view count_type,
+                            std::string_view query_type) {
+  return std::visit(
+      [&](const auto &ff) -> py::object {
+        return file_fetch_sum(ff, range1, range2, normalization, count_type, query_type);
+      },
+      f.get());
+}
+
+inline std::int64_t fetch_nnz(const hictk::File &f, std::string_view range1,
+                              std::string_view range2, std::string_view query_type) {
+  return std::visit([&](const auto &ff) { return file_fetch_nnz(ff, range1, range2, query_type); },
+                    f.get());
+}
+
 }  // namespace hictkpy::file
diff --git a/src/hictkpy_hic.hpp b/src/hictkpy_hic.hpp
@@ -37,4 +37,15 @@ inline py::object fetch_dense(const hictk::hic::File &f, std::string_view range1
                               std::string_view count_type, std::string_view query_type) {
   return file_fetch_dense(f, range1, range2, normalization, count_type, query_type);
 }
+
+inline py::object fetch_sum(const hictk::hic::File &f, std::string_view range1,
+                            std::string_view range2, std::string_view normalization,
+                            std::string_view count_type, std::string_view query_type) {
+  return file_fetch_sum(f, range1, range2, normalization, count_type, query_type);
+}
+
+inline std::int64_t fetch_nnz(const hictk::hic::File &f, std::string_view range1,
+                              std::string_view range2, std::string_view query_type) {
+  return file_fetch_nnz(f, range1, range2, query_type);
+}
 }  // namespace hictkpy::hic