Merge pull request #128 from paulsengroup/feat/file-writer-bins

Add accessors for BinTable objects to FileWriter classes

docs/api/cooler.rst

@@ -23,6 +23,7 @@ Cooler API
 
    .. automethod:: __init__
    .. automethod:: add_pixels
+   .. automethod:: bins
    .. automethod:: chromosomes
    .. automethod:: finalize
    .. automethod:: path

docs/api/hic.rst

@@ -12,6 +12,7 @@ Hi-C API
 
    .. automethod:: __init__
    .. automethod:: add_pixels
+   .. automethod:: bins
    .. automethod:: chromosomes
    .. automethod:: finalize
    .. automethod:: path

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"),

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"),

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);
 

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);
 

test/test_file_creation_cool.py

@@ -44,6 +44,7 @@ def test_accessors(self, file, resolution, tmpdir):
         else:
             assert w.resolution() == resolution
         assert w.chromosomes() == bins.chromosomes()
+        assert len(w.bins().to_df().compare(bins.to_df())) == 0
 
     def test_file_creation_thin_pixel(self, file, resolution, tmpdir):
         f = hictkpy.File(file, resolution)

test/test_file_creation_hic.py

@@ -31,6 +31,20 @@ def setup_method():
         logging.basicConfig(level="INFO", force=True)
         logging.getLogger().setLevel("INFO")
 
+    def test_accessors(self, file, resolution, tmpdir):
+        bins = hictkpy.File(file, resolution).bins()
+        if bins.type() != "fixed":
+            pytest.skip(f'BinTable of file "{file}" does not have fixed bins.')
+
+        path = tmpdir / "test.hic"
+        w = hictkpy.hic.FileWriter(path, bins)
+
+        assert str(w).startswith("HiCFileWriter(")
+        assert w.path() == path
+        assert w.resolutions() == [resolution]
+        assert w.chromosomes() == bins.chromosomes()
+        assert len(w.bins(resolution).to_df().compare(bins.to_df())) == 0
+
     def test_file_creation_thin_pixel(self, file, resolution, tmpdir):
         f = hictkpy.File(file, resolution)
         if f.bins().type() != "fixed":
@@ -39,7 +53,7 @@ def test_file_creation_thin_pixel(self, file, resolution, tmpdir):
         df = f.fetch(join=False).to_df()
         expected_sum = df["count"].sum()
 
-        path = tmpdir / "test1.hic"
+        path = tmpdir / "test.hic"
         w = hictkpy.hic.FileWriter(path, f.chromosomes(), f.resolution())
 
         chunk_size = 1000
@@ -67,7 +81,7 @@ def test_file_creation(self, file, resolution, tmpdir):
         df = f.fetch(join=True).to_df()
         expected_sum = df["count"].sum()
 
-        path = tmpdir / "test2.hic"
+        path = tmpdir / "test.hic"
         w = hictkpy.hic.FileWriter(path, f.chromosomes(), f.resolution())
 
         chunk_size = 1000
@@ -93,7 +107,7 @@ def test_file_creation_bin_table(self, file, resolution, tmpdir):
         df = f.fetch(join=True).to_df()
         expected_sum = df["count"].sum()
 
-        path = tmpdir / "test2.hic"
+        path = tmpdir / "test.hic"
         if f.bins().type() != "fixed":
             with pytest.raises(Exception):
                 hictkpy.hic.FileWriter(path, f.bins())