Merge pull request #57 from paulsengroup/bump/nanobind

Bump nanobind and support stubs

.gitignore

@@ -167,3 +167,6 @@ cython_debug/
 # and can be added to the global gitignore or merged into this file. For a more nuclear
 # option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
+
+py.typed
+*.pyi

CMakeLists.txt

@@ -34,8 +34,8 @@ FetchContent_Declare(
 
 FetchContent_Declare(
  nanobind
- URL "${CMAKE_CURRENT_SOURCE_DIR}/external/nanobind-v1.9.2.tar.xz"
- URL_HASH "SHA256=2f511306aea0b39c1915b85db77cf4d861056dec528fb56a3e7a5f79fdae01c4"
+ URL "${CMAKE_CURRENT_SOURCE_DIR}/external/nanobind-v2.0.0.tar.xz"
+ URL_HASH "SHA256=e9accc281605cdab732987756421244ab8b55e026c0fdbae93fa90d95f4dce38"
  EXCLUDE_FROM_ALL
  SYSTEM)
 

pyproject.toml

@@ -6,6 +6,10 @@
 requires = [
  "conan>=2.0.5",
  "scikit-build-core>=0.8",
+ "numpy",
+ "pandas>=2.1.0,!=2.2.0",
+ "scipy",
+ "typing_extensions",
 ]
 
 build-backend = "scikit_build_core.build"
@@ -66,7 +70,7 @@ filterwarnings = [
 ]
 
 [tool.cibuildwheel]
-skip = ["*musllinux*"]
+skip = ["*musllinux*", "pp*"]
 test-command = "python -m pytest {project}/test"
 test-extras = ["test"]
 test-skip = ["*universal2", "pp*"]

src/CMakeLists.txt

@@ -29,3 +29,22 @@ target_link_libraries(
  hictk::hic)
 
 install(TARGETS _hictkpy LIBRARY DESTINATION hictkpy)
+
+# Automatic stub generation for hictkpy does not work, instead run:
+# venv/bin/python -m nanobind.stubgen -m hictkpy -o src/hictkpy/hictkpy.pyi
+# Remember to remove the version annotation from the stub
+
+nanobind_add_stub(
+ hictkpy._hictkpy_stub
+ INSTALL_TIME
+ MODULE
+ hictkpy._hictkpy
+ MARKER_FILE
+ "${CMAKE_CURRENT_BINARY_DIR}/hictkpy/py.typed"
+ OUTPUT
+ "${CMAKE_CURRENT_BINARY_DIR}/hictkpy/_hictkpy.pyi"
+ VERBOSE)
+
+install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/hictkpy/hictkpy.pyi" DESTINATION hictkpy)
+install(FILES "${CMAKE_CURRENT_BINARY_DIR}/hictkpy/_hictkpy.pyi" DESTINATION hictkpy)
+install(FILES "${CMAKE_CURRENT_BINARY_DIR}/hictkpy/py.typed" DESTINATION hictkpy)

src/hictkpy.cpp

@@ -74,27 +74,35 @@ static void declare_pixel_selector_class(nb::module_ &m) {
  m, "PixelSelector",
  "Class representing pixels overlapping with the given genomic intervals.");
 
- sel.def(nb::init<std::shared_ptr<const hictk::cooler::PixelSelector>, std::string_view, bool, bool>(),
- nb::arg("selector"), nb::arg("type"), nb::arg("join"), nb::arg("_mirror"));
- sel.def(nb::init<std::shared_ptr<const hictk::hic::PixelSelector>, std::string_view, bool, bool>(),
- nb::arg("selector"), nb::arg("type"), nb::arg("join"), nb::arg("_mirror"));
- sel.def(nb::init<std::shared_ptr<const hictk::hic::PixelSelectorAll>, std::string_view, bool, bool>(),
- nb::arg("selector"), nb::arg("type"), nb::arg("join"), nb::arg("_mirror"));
+ sel.def(
+ nb::init<std::shared_ptr<const hictk::cooler::PixelSelector>, std::string_view, bool, bool>(),
+ nb::arg("selector"), nb::arg("type"), nb::arg("join"), nb::arg("_mirror"));
+ sel.def(
+ nb::init<std::shared_ptr<const hictk::hic::PixelSelector>, std::string_view, bool, bool>(),
+ nb::arg("selector"), nb::arg("type"), nb::arg("join"), nb::arg("_mirror"));
+ sel.def(
+ nb::init<std::shared_ptr<const hictk::hic::PixelSelectorAll>, std::string_view, bool, bool>(),
+ nb::arg("selector"), nb::arg("type"), nb::arg("join"), nb::arg("_mirror"));
 
  sel.def("__repr__", &PixelSelector::repr);
 
  sel.def("coord1", &PixelSelector::get_coord1, "Get query coordinates for the first dimension.");
  sel.def("coord2", &PixelSelector::get_coord2, "Get query coordinates for the second dimension.");
 
- sel.def("__iter__", &PixelSelector::make_iterable, nb::keep_alive<0, 1>());
+ sel.def("__iter__", &PixelSelector::make_iterable, nb::keep_alive<0, 1>(),
+ nb::sig("def __iter__(self) -> ThinPixelInt | ThinPixelFP | PixelInt | PixelFP"),
+ "Return an iterator over the selected pixels.");
 
- sel.def("to_df", &PixelSelector::to_df, "Retrieve interactions as a pandas DataFrame.");
- sel.def("to_numpy", &PixelSelector::to_numpy, "Retrieve interactions as a numpy 2D matrix.");
- sel.def("to_coo", &PixelSelector::to_coo, "Retrieve interactions as a scipy.sparse.coo_matrix.");
+ sel.def("to_df", &PixelSelector::to_df, nb::sig("def to_df(self) -> pandas.DataFrame"),
+ "Retrieve interactions as a pandas DataFrame.");
+ sel.def("to_numpy", &PixelSelector::to_numpy, nb::sig("def to_numpy(self) -> numpy.ndarray"),
+ "Retrieve interactions as a numpy 2D matrix.");
+ sel.def("to_coo", &PixelSelector::to_coo, nb::sig("def to_coo(self) -> scipy.sparse.coo_matrix"),
+ "Retrieve interactions as a SciPy COO matrix.");
 
  sel.def("nnz", &PixelSelector::nnz,
  "Get the number of non-zero entries for the current pixel selection.");
- sel.def("sum", &PixelSelector::sum,
+ sel.def("sum", &PixelSelector::sum, nb::sig("def sum(self) -> int | float"),
  "Get the total number of interactions for the current pixel selection.");
 }
 
@@ -118,7 +126,8 @@ static void declare_file_class(nb::module_ &m) {
 
  file.def("chromosomes", &get_chromosomes_from_file<hictk::File>, nb::arg("include_all") = false,
  "Get chromosomes sizes as a dictionary mapping names to sizes.");
- file.def("bins", &get_bins_from_file<hictk::File>, "Get bins as a pandas DataFrame.");
+ file.def("bins", &get_bins_from_file<hictk::File>, nb::sig("def bins(self) -> pandas.DataFrame"),
+ "Get bins as a pandas DataFrame.");
 
  file.def("resolution", &hictk::File::resolution, "Get the bin size in bp.");
  file.def("nbins", &hictk::File::nbins, "Get the total number of bins.");
@@ -154,7 +163,7 @@ static void declare_multires_file_class(nb::module_ &m) {
  mres_file.def("resolutions", &hictk::MultiResFile::resolutions,
  "Get the list of available resolutions.");
  mres_file.def("__getitem__", &hictk::MultiResFile::open,
- "Open the Cooler file corresponding to the resolution given as input.");
+ "Open the Cooler or .hic file corresponding to the resolution given as input.");
 }
 
 static void declare_singlecell_file_class(nb::module_ &m) {
@@ -175,7 +184,7 @@ static void declare_singlecell_file_class(nb::module_ &m) {
  nb::arg("include_all") = false,
  "Get chromosomes sizes as a dictionary mapping names to sizes.");
  scell_file.def("bins", &get_bins_from_file<hictk::cooler::SingleCellFile>,
- "Get bins as a pandas DataFrame.");
+ nb::sig("def bins(self) -> pandas.DataFrame"), "Get bins as a pandas DataFrame.");
  scell_file.def("attributes", &singlecell_file::get_attrs, "Get file attributes as a dictionary.");
  scell_file.def("cells", &singlecell_file::get_cells, "Get the list of available cells.");
  scell_file.def("__getitem__", &singlecell_file::getitem,
@@ -211,7 +220,8 @@ static void declare_hic_file_writer_class(nb::module_ &m) {
  nb::arg("include_all") = false,
  "Get chromosomes sizes as a dictionary mapping names to sizes.");
 
- writer.def("add_pixels", &hictkpy::HiCFileWriter::add_pixels, nb::arg("pixels"),
+ writer.def("add_pixels", &hictkpy::HiCFileWriter::add_pixels,
+ nb::sig("def add_pixels(self, pixels: pd.DataFrame) -> None"), nb::arg("pixels"),
  "Add pixels from a pandas DataFrame containing pixels in COO or BG2 format (i.e. "
  "either with columns=[bin1_id, bin2_id, count] or with columns=[chrom1, start1, end1, "
  "chrom2, start2, end2, count].");
@@ -238,7 +248,8 @@ static void declare_cooler_file_writer_class(nb::module_ &m) {
  nb::arg("include_all") = false,
  "Get chromosomes sizes as a dictionary mapping names to sizes.");
 
- writer.def("add_pixels", &hictkpy::CoolFileWriter::add_pixels, nb::arg("pixels"),
+ writer.def("add_pixels", &hictkpy::CoolFileWriter::add_pixels,
+ nb::sig("def add_pixels(self, pixels: pandas.DataFrame)"), nb::arg("pixels"),
  "Add pixels from a pandas DataFrame containing pixels in COO or BG2 format (i.e. "
  "either with columns=[bin1_id, bin2_id, count] or with columns=[chrom1, start1, end1, "
  "chrom2, start2, end2, count].");

src/hictkpy/hictkpy.pyi

@@ -0,0 +1,22 @@
+# Copyright (C) 2024 Roberto Rossini <[email protected]>
+#
+# SPDX-License-Identifier: MIT
+
+from . import (
+ _hictkpy.cooler as _hictkpy.cooler,
+ _hictkpy.hic as _hictkpy.hic
+)
+from ._hictkpy import (
+ File as File,
+ MultiResFile as MultiResFile,
+ PixelSelector as PixelSelector,
+ is_cooler as is_cooler,
+ is_hic as is_hic,
+ is_mcool_file as is_mcool_file,
+ is_scool_file as is_scool_file
+)
+
+
+__all__: list = ['__doc__', 'File', 'MultiResFile', 'PixelSelector', 'is_cooler', 'is_mcool_file', 'is_scool_file', 'is_hic', 'cooler', 'hic', '__hictk_version__']
+
+__hictk_version__: str

src/hictkpy_file_creation.cpp

@@ -23,8 +23,8 @@ namespace hictkpy {
 
 template <typename N>
 std::vector<hictk::ThinPixel<N>> coo_df_to_thin_pixels(nanobind::object df, bool sorted) {
- using BufferT1 = nb::ndarray<nb::numpy, nb::shape<nb::any>, std::uint64_t>;
- using BufferT2 = nb::ndarray<nb::numpy, nb::shape<nb::any>, N>;
+ using BufferT1 = nb::ndarray<nb::numpy, nb::shape<-1>, std::uint64_t>;
+ using BufferT2 = nb::ndarray<nb::numpy, nb::shape<-1>, N>;
 
  auto bin1_ids_np = nb::cast<BufferT1>(df.attr("__getitem__")("bin1_id").attr("to_numpy")());
  auto bin2_ids_np = nb::cast<BufferT1>(df.attr("__getitem__")("bin2_id").attr("to_numpy")());
@@ -49,8 +49,8 @@ std::vector<hictk::ThinPixel<N>> coo_df_to_thin_pixels(nanobind::object df, bool
 template <typename N>
 std::vector<hictk::ThinPixel<N>> bg2_df_to_thin_pixels(const hictk::BinTable &bin_table,
  nanobind::object df, bool sorted) {
- using BufferT1 = nb::ndarray<nb::numpy, nb::shape<nb::any>, std::uint32_t>;
- using BufferT2 = nb::ndarray<nb::numpy, nb::shape<nb::any>, N>;
+ using BufferT1 = nb::ndarray<nb::numpy, nb::shape<-1>, std::uint32_t>;
+ using BufferT2 = nb::ndarray<nb::numpy, nb::shape<-1>, N>;
 
  auto chrom1 = nb::cast<nb::list>(df.attr("__getitem__")("chrom1").attr("tolist")());
  auto start1_np = nb::cast<BufferT1>(df.attr("__getitem__")("start1").attr("to_numpy")());

src/hictkpy_pixel_selector.cpp

@@ -114,42 +114,42 @@ auto PixelSelector::get_coord2() const -> PixelCoordTuple {
  c.bin2.chrom().name(), c.bin2.start(), c.bin2.end())};
 }
 
-nb::iterator PixelSelector::make_iterable() const {
+nb::object PixelSelector::make_iterable() const {
  if (mirror) {
  throw std::runtime_error(
  "iterating through the pixels for a query overlapping the lower triangle is not supported");
  }
 
  if (join) {
  return std::visit(
- [&](const auto& s) {
+ [&](const auto& s) -> nb::object {
  if (int_pixels()) {
  using T = std::int32_t;
  auto jsel = hictk::transformers::JoinGenomicCoords(
  s->template begin<T>(), s->template end<T>(),
  std::make_shared<const hictk::BinTable>(bins()));
- return nb::make_iterator(nb::type<PixelSelector>(), "PixelIterator", jsel.begin(),
- jsel.end());
+ return nb::cast(nb::make_iterator(nb::type<PixelSelector>(), "PixelIterator",
+  jsel.begin(), jsel.end()));
  }
  using T = double;
  auto jsel = hictk::transformers::JoinGenomicCoords(
  s->template begin<T>(), s->template end<T>(),
  std::make_shared<const hictk::BinTable>(bins()));
- return nb::make_iterator(nb::type<PixelSelector>(), "PixelIterator", jsel.begin(),
- jsel.end());
+ return nb::cast(nb::make_iterator(nb::type<PixelSelector>(), "PixelIterator",
+  jsel.begin(), jsel.end()));
  },
  selector);
  }
  return std::visit(
- [&](const auto& s) {
+ [&](const auto& s) -> nb::object {
  if (int_pixels()) {
  using T = std::int32_t;
- return nb::make_iterator(nb::type<PixelSelector>(), "PixelIterator",
- s->template begin<T>(), s->template end<T>());
+ return nb::cast(nb::make_iterator(nb::type<PixelSelector>(), "PixelIterator",
+  s->template begin<T>(), s->template end<T>()));
  }
  using T = double;
- return nb::make_iterator(nb::type<PixelSelector>(), "PixelIterator", s->template begin<T>(),
- s->template end<T>());
+ return nb::cast(nb::make_iterator(nb::type<PixelSelector>(), "PixelIterator",
+  s->template begin<T>(), s->template end<T>()));
  },
  selector);
 }

src/include/hictkpy/common.hpp

@@ -109,7 +109,7 @@ template <typename T>
 template <typename T>
 struct Dynamic1DA {
  private:
- using BufferT = nb::ndarray<nb::numpy, nb::shape<nb::any>, T>;
+ using BufferT = nb::ndarray<nb::numpy, nb::shape<-1>, T>;
  using VectorT = decltype(std::declval<BufferT>().view());
  nb::object _np_array{};
  BufferT _buff{};
@@ -169,8 +169,9 @@ struct Dynamic1DA {
 };
 
 template <typename File>
-inline nb::dict get_chromosomes_from_file(const File &f, bool include_all = false) {
- nb::dict py_chroms{}; // NOLINT
+inline nb::typed<nb::dict, std::string, std::uint32_t> get_chromosomes_from_file(
+ const File &f, bool include_all = false) {
+ nb::typed<nb::dict, std::string, std::uint32_t> py_chroms{}; // NOLINT
  for (const auto &chrom : f.chromosomes()) {
  if (!include_all && chrom.is_all()) {
  continue;
@@ -301,7 +302,7 @@ inline nb::object pixel_iterators_to_numpy(PixelIt first_pixel, PixelIt last_pix
  const auto dtype = np.attr("dtype")(map_type_to_dtype<N>());
  auto buffer = np.attr("zeros")(std::vector<std::int64_t>{num_rows, num_cols}, "dtype"_a = dtype);
 
- using Shape = nb::shape<nb::any, nb::any>;
+ using Shape = nb::shape<-1, -1>;
  using MatrixT = nb::ndarray<nb::numpy, N, Shape>;
 
  auto matrix = nb::cast<MatrixT>(buffer);

src/include/hictkpy/pixel_selector.hpp

@@ -6,6 +6,7 @@
 
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/string_view.h>
+#include <nanobind/stl/tuple.h>
 
 #include <cstdint>
 #include <memory>
@@ -50,7 +51,7 @@ struct PixelSelector {
  [[nodiscard]] auto get_coord1() const -> PixelCoordTuple;
  [[nodiscard]] auto get_coord2() const -> PixelCoordTuple;
 
- [[nodiscard]] nanobind::iterator make_iterable() const;
+ [[nodiscard]] nanobind::object make_iterable() const;
  [[nodiscard]] nanobind::object to_df() const;
  [[nodiscard]] nanobind::object to_coo() const;
  [[nodiscard]] nanobind::object to_numpy() const;