Skip to content

Add individual nodes method #3153

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
benjeffery merged 3 commits into from
May 7, 2025
+263 −50
Merged

Add individual nodes method #3153

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
251ad37
Fix CI linting not finding clang
benjeffery May 7, 2025
9bf6a7f
Add individuals nodes method
benjeffery May 3, 2025
0ab2825
Update dev version
benjeffery May 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions .github/workflows/tests.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,10 @@ jobs:
- name: install clang-format
if: steps.clang_format.outputs.cache-hit != 'true'
run: |
python -m venv env
source env/bin/activate
pip install clang-format==6.0.1
pip install setuptools clang-format==6.0.1
clang-format -version
sudo cp /opt/hostedtoolcache/Python/3.12.10/x64/bin/clang-format /opt/hostedtoolcache/Python/3.12.10/x64/bin/clang-format-6.0
clang-format-6.0 -version
- uses: pre-commit/action@v3.0.1

benchmark:
Expand Down
2 changes: 1 addition & 1 deletion c/tests/test_core.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -532,7 +532,7 @@ test_bit_arrays(void)
// NB: This test is only valid for the 32 bit implementation of bit arrays. If we
// were to change the chunk size of a bit array, we'd need to update these tests
tsk_bit_array_t arr;
tsk_id_t items_truth[64] = {0}, items[64] = {0};
tsk_id_t items_truth[64] = { 0 }, items[64] = { 0 };
tsk_size_t n_items = 0, n_items_truth = 0;

// test item retrieval
Expand Down
6 changes: 3 additions & 3 deletions c/tests/test_stats.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -348,7 +348,7 @@ verify_pair_coalescence_counts(tsk_treeseq_t *ts, tsk_flags_t options)

/* test various bin assignments */
for (i = 0; i < N; i++) {
node_bin_map[i] = ((tsk_id_t) (i % B));
node_bin_map[i] = ((tsk_id_t)(i % B));
}
ret = tsk_treeseq_pair_coalescence_counts(ts, P, sample_set_sizes, sample_sets, I,
index_tuples, T, breakpoints, B, node_bin_map, options, C_B);
Expand Down Expand Up @@ -3694,8 +3694,8 @@ static void
test_pair_coalescence_counts_missing(void)
{
tsk_treeseq_t ts;
tsk_treeseq_from_text(&ts, 5, missing_ex_nodes, missing_ex_edges, NULL,
NULL, NULL, NULL, NULL, 0);
tsk_treeseq_from_text(
&ts, 5, missing_ex_nodes, missing_ex_edges, NULL, NULL, NULL, NULL, NULL, 0);
verify_pair_coalescence_counts(&ts, 0);
verify_pair_coalescence_counts(&ts, TSK_STAT_SPAN_NORMALISE);
tsk_treeseq_free(&ts);
Expand Down
46 changes: 22 additions & 24 deletions c/tests/testlib.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -373,31 +373,29 @@ const char *empty_ex_nodes = "1 0.0 0 -1\n"
const char *empty_ex_edges = "";

/*** An example of a tree sequence with missing marginal trees. ***/
/*
| 4 | | 4 |
| / \ | | / \ |
| 3 \ | | / 3 |
| / \ \ | | / / \ |
| 0 1 2 | | 0 1 2 |
|-|-----------|-|-----------|-|
0 1 2 3 4 5
/*
| 4 | | 4 |
| / \ | | / \ |
| 3 \ | | / 3 |
| / \ \ | | / / \ |
| 0 1 2 | | 0 1 2 |
|-|-----------|-|-----------|-|
0 1 2 3 4 5
*/
const char *missing_ex_nodes =
"1 0.0 0 -1\n"
"1 0.0 0 -1\n"
"1 0.0 0 -1\n"
"0 1.0 0 -1\n"
"0 2.0 0 -1\n";

const char *missing_ex_edges =
"1.0 2.0 3 0\n"
"1.0 2.0 3 1\n"
"3.0 4.0 3 1\n"
"3.0 4.0 3 2\n"
"3.0 4.0 4 0\n"
"1.0 2.0 4 2\n"
"1.0 2.0 4 3\n"
"3.0 4.0 4 3\n";
const char *missing_ex_nodes = "1 0.0 0 -1\n"
"1 0.0 0 -1\n"
"1 0.0 0 -1\n"
"0 1.0 0 -1\n"
"0 2.0 0 -1\n";

const char *missing_ex_edges = "1.0 2.0 3 0\n"
"1.0 2.0 3 1\n"
"3.0 4.0 3 1\n"
"3.0 4.0 3 2\n"
"3.0 4.0 4 0\n"
"1.0 2.0 4 2\n"
"1.0 2.0 4 3\n"
"3.0 4.0 4 3\n";

/* Simple utilities to parse text so we can write declaritive
* tests. This is not intended as a robust general input mechanism.
Expand Down
3 changes: 2 additions & 1 deletion c/tskit/core.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,8 @@ tsk_strerror_internal(int err)
break;
case TSK_ERR_FILE_VERSION_TOO_OLD:
ret = "tskit file version too old. Please upgrade using the "
"'tskit upgrade' command from tskit version<0.6.2. (TSK_ERR_FILE_VERSION_TOO_OLD)";
"'tskit upgrade' command from tskit version<0.6.2. "
"(TSK_ERR_FILE_VERSION_TOO_OLD)";
break;
case TSK_ERR_FILE_VERSION_TOO_NEW:
ret = "tskit file version is too new for this instance. "
Expand Down
10 changes: 10 additions & 0 deletions python/CHANGELOG.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,13 @@
--------------------
[0.6.4] - 2025-XX-XX
--------------------

**Features**

- Add ``TreeSequence.individuals_nodes`` attribute to return the nodes
associated with each individual as a numpy array.
(:user:`benjeffery`, :pr:`3153`)

--------------------
[0.6.3] - 2025-04-28
--------------------
Expand Down
104 changes: 87 additions & 17 deletions python/_tskitmodule.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,8 @@
const char *not_kas_format_msg
= "File not in kastore format. Either the file is corrupt or it is not a "
"tskit tree sequence file. It may be a legacy HDF file upgradable with "
"`tskit upgrade` from tskit version<0.6.2 or a compressed tree sequence file that can be decompressed "
"`tskit upgrade` from tskit version<0.6.2 or a compressed tree sequence file "
"that can be decompressed "
"with `tszip`.";
const char *ibd_pairs_not_stored_msg
= "Sample pairs are not stored by default "
Expand Down Expand Up @@ -8724,6 +8725,61 @@
return ret;
}

static PyObject *
TreeSequence_get_individuals_nodes(TreeSequence *self)
{
PyObject *ret = NULL;
PyArrayObject *ret_array = NULL;
npy_intp dims[2];
tsk_size_t ploidy;
tsk_size_t max_ploidy = 0;
tsk_id_t *node_mem;
tsk_size_t j;
tsk_size_t num_individuals;
tsk_id_t *const *individual_nodes;
const tsk_size_t *individual_nodes_length;

if (TreeSequence_check_state(self) != 0) {
goto out;
}

num_individuals = tsk_treeseq_get_num_individuals(self->tree_sequence);
individual_nodes = self->tree_sequence->individual_nodes;
individual_nodes_length = self->tree_sequence->individual_nodes_length;

for (tsk_id_t i = 0; i < (tsk_id_t) num_individuals; i++) {
ploidy = individual_nodes_length[i];
if (ploidy > max_ploidy) {
max_ploidy = ploidy;
}
}

dims[0] = (npy_intp) num_individuals;
dims[1] = (npy_intp) max_ploidy;
ret_array = (PyArrayObject *) PyArray_SimpleNew(2, dims, NPY_INT32);
if (ret_array == NULL) {
goto out;

Check warning on line 8761 in python/_tskitmodule.c

View check run for this annotation

Codecov / codecov/patch

python/_tskitmodule.c#L8761 

Added line #L8761 was not covered by tests
}

/* Fill with -1 (TSK_NULL) */
node_mem = (tsk_id_t *) PyArray_DATA(ret_array);
memset(node_mem, 0xFF, PyArray_NBYTES(ret_array));

for (tsk_id_t i = 0; i < (tsk_id_t) num_individuals; i++) {
ploidy = individual_nodes_length[i];
for (j = 0; j < ploidy; j++) {
node_mem[i * max_ploidy + j] = individual_nodes[i][j];
}
}

ret = (PyObject *) ret_array;
ret_array = NULL;

out:
Py_XDECREF(ret_array);
return ret;
}

static PyObject *
TreeSequence_genealogical_nearest_neighbours(
TreeSequence *self, PyObject *args, PyObject *kwds)
Expand Down Expand Up @@ -9715,12 +9771,12 @@
if (result_array == NULL) {
goto out;
}
Py_BEGIN_ALLOW_THREADS
err = method(self->tree_sequence, w_shape[1], PyArray_DATA(weights_array),
num_windows, PyArray_DATA(windows_array), num_focal_nodes,
PyArray_DATA(focal_nodes_array), PyArray_DATA(result_array), options);
Py_END_ALLOW_THREADS
if (err != 0) {
Py_BEGIN_ALLOW_THREADS err
= method(self->tree_sequence, w_shape[1], PyArray_DATA(weights_array),
num_windows, PyArray_DATA(windows_array), num_focal_nodes,
PyArray_DATA(focal_nodes_array), PyArray_DATA(result_array), options);
Py_END_ALLOW_THREADS if (err != 0)
{
handle_library_error(err);
goto out;
}
Expand Down Expand Up @@ -10744,7 +10800,8 @@
goto out;
}
individuals = self->tree_sequence->tables->individuals;
ret = TreeSequence_make_array(self, individuals.metadata_length, NPY_UINT8, individuals.metadata);
ret = TreeSequence_make_array(
self, individuals.metadata_length, NPY_UINT8, individuals.metadata);
out:
return ret;
}
Expand Down Expand Up @@ -10835,7 +10892,8 @@
goto out;
}
nodes = self->tree_sequence->tables->nodes;
ret = TreeSequence_make_array(self, nodes.metadata_length, NPY_UINT8, nodes.metadata);
ret = TreeSequence_make_array(
self, nodes.metadata_length, NPY_UINT8, nodes.metadata);
out:
return ret;
}
Expand Down Expand Up @@ -10926,7 +10984,8 @@
goto out;
}
edges = self->tree_sequence->tables->edges;
ret = TreeSequence_make_array(self, edges.metadata_length, NPY_UINT8, edges.metadata);
ret = TreeSequence_make_array(
self, edges.metadata_length, NPY_UINT8, edges.metadata);
out:
return ret;
}
Expand All @@ -10941,7 +11000,8 @@
goto out;
}
edges = self->tree_sequence->tables->edges;
ret = TreeSequence_make_array(self, edges.num_rows + 1, NPY_UINT64, edges.metadata_offset);
ret = TreeSequence_make_array(
self, edges.num_rows + 1, NPY_UINT64, edges.metadata_offset);
out:
return ret;
}
Expand Down Expand Up @@ -10971,7 +11031,8 @@
goto out;
}
sites = self->tree_sequence->tables->sites;
ret = TreeSequence_make_array(self, sites.metadata_length, NPY_UINT8, sites.metadata);
ret = TreeSequence_make_array(
self, sites.metadata_length, NPY_UINT8, sites.metadata);
out:
return ret;
}
Expand All @@ -10986,7 +11047,8 @@
goto out;
}
sites = self->tree_sequence->tables->sites;
ret = TreeSequence_make_array(self, sites.num_rows + 1, NPY_UINT64, sites.metadata_offset);
ret = TreeSequence_make_array(
self, sites.num_rows + 1, NPY_UINT64, sites.metadata_offset);
out:
return ret;
}
Expand Down Expand Up @@ -11061,7 +11123,8 @@
goto out;
}
mutations = self->tree_sequence->tables->mutations;
ret = TreeSequence_make_array(self, mutations.metadata_length, NPY_UINT8, mutations.metadata);
ret = TreeSequence_make_array(
self, mutations.metadata_length, NPY_UINT8, mutations.metadata);
out:
return ret;
}
Expand All @@ -11076,7 +11139,8 @@
goto out;
}
mutations = self->tree_sequence->tables->mutations;
ret = TreeSequence_make_array(self, mutations.num_rows + 1, NPY_UINT64, mutations.metadata_offset);
ret = TreeSequence_make_array(
self, mutations.num_rows + 1, NPY_UINT64, mutations.metadata_offset);
out:
return ret;
}
Expand Down Expand Up @@ -11185,7 +11249,8 @@
goto out;
}
migrations = self->tree_sequence->tables->migrations;
ret = TreeSequence_make_array(self, migrations.metadata_length, NPY_UINT8, migrations.metadata);
ret = TreeSequence_make_array(
self, migrations.metadata_length, NPY_UINT8, migrations.metadata);
out:
return ret;
}
Expand Down Expand Up @@ -11216,7 +11281,8 @@
goto out;
}
populations = self->tree_sequence->tables->populations;
ret = TreeSequence_make_array(self, populations.metadata_length, NPY_UINT8, populations.metadata);
ret = TreeSequence_make_array(
self, populations.metadata_length, NPY_UINT8, populations.metadata);
out:
return ret;
}
Expand Down Expand Up @@ -11414,6 +11480,10 @@
.ml_meth = (PyCFunction) TreeSequence_get_individuals_time,
.ml_flags = METH_NOARGS,
.ml_doc = "Returns the vector of per-individual times." },
{ .ml_name = "get_individuals_nodes",
.ml_meth = (PyCFunction) TreeSequence_get_individuals_nodes,
.ml_flags = METH_NOARGS,
.ml_doc = "Returns an array of the node ids for each individual" },
{ .ml_name = "genealogical_nearest_neighbours",
.ml_meth = (PyCFunction) TreeSequence_genealogical_nearest_neighbours,
.ml_flags = METH_VARARGS | METH_KEYWORDS,
Expand Down
Loading
Loading