Add accessors for BinTable objects to FileWriter classes

paulsengroup · Nov 11, 2024 · eb802e7 · eb802e7
1 parent 6b2f865
commit eb802e7
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 0 deletions.
diff --git a/src/cooler_file_writer.cpp b/src/cooler_file_writer.cpp
@@ -71,6 +71,14 @@ const hictk::Reference &CoolerFileWriter::chromosomes() const {
   return ref;
 }
 
+std::shared_ptr<const hictk::BinTable> CoolerFileWriter::bins_ptr() const noexcept {
+  if (!_w) {
+    return {};
+  }
+
+  return _w->bins_ptr();
+}
+
 void CoolerFileWriter::add_pixels(const nb::object &df) {
   if (!_w.has_value()) {
     throw std::runtime_error(
@@ -193,6 +201,8 @@ void CoolerFileWriter::bind(nb::module_ &m) {
              nb::arg("include_ALL") = false,
              "Get chromosomes sizes as a dictionary mapping names to sizes.",
              nb::rv_policy::take_ownership);
+  writer.def("bins", &get_bins_from_object<hictkpy::CoolerFileWriter>, "Get table of bins.",
+             nb::sig("def bins(self) -> hictkpy.BinTable"), nb::rv_policy::move);
 
   writer.def("add_pixels", &hictkpy::CoolerFileWriter::add_pixels,
              nb::sig("def add_pixels(self, pixels: pandas.DataFrame)"), nb::arg("pixels"),

diff --git a/src/hic_file_writer.cpp b/src/hic_file_writer.cpp
@@ -105,6 +105,10 @@ const std::vector<std::uint32_t> &HiCFileWriter::resolutions() const noexcept {
 
 const hictk::Reference &HiCFileWriter::chromosomes() const { return _w.chromosomes(); }
 
+hictkpy::BinTable HiCFileWriter::bins(std::uint32_t resolution) const {
+  return hictkpy::BinTable{_w.bins(resolution)};
+}
+
 void HiCFileWriter::add_pixels(const nb::object &df) {
   if (_finalized) {
     throw std::runtime_error(
@@ -172,6 +176,8 @@ void HiCFileWriter::bind(nb::module_ &m) {
              nb::arg("include_ALL") = false,
              "Get chromosomes sizes as a dictionary mapping names to sizes.",
              nb::rv_policy::take_ownership);
+  writer.def("bins", &hictkpy::HiCFileWriter::bins, "Get table of bins for the given resolution.",
+             nb::sig("def bins(self, resolution: int) -> hictkpy.BinTable"), nb::rv_policy::move);
 
   writer.def("add_pixels", &hictkpy::HiCFileWriter::add_pixels,
              nb::sig("def add_pixels(self, pixels: pd.DataFrame) -> None"), nb::arg("pixels"),

diff --git a/src/include/hictkpy/cooler_file_writer.hpp b/src/include/hictkpy/cooler_file_writer.hpp
@@ -39,6 +39,7 @@ class CoolerFileWriter {
   [[nodiscard]] std::uint32_t resolution() const noexcept;
 
   [[nodiscard]] const hictk::Reference& chromosomes() const;
+  [[nodiscard]] std::shared_ptr<const hictk::BinTable> bins_ptr() const noexcept;
 
   void add_pixels(const nanobind::object& df);
 

diff --git a/src/include/hictkpy/hic_file_writer.hpp b/src/include/hictkpy/hic_file_writer.hpp
@@ -43,6 +43,7 @@ class HiCFileWriter {
   [[nodiscard]] const std::vector<std::uint32_t>& resolutions() const noexcept;
 
   [[nodiscard]] const hictk::Reference& chromosomes() const;
+  [[nodiscard]] hictkpy::BinTable bins(std::uint32_t resolution) const;
 
   void add_pixels(const nanobind::object& df);