tskit 1.0.2__tar.gz → 1.0.3__tar.gz

{tskit-1.0.2/tskit.egg-info → tskit-1.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tskit
-Version: 1.0.2
+Version: 1.0.3
 Summary: The tree sequence toolkit.
 Author-email: Tskit Developers <admin@tskit.dev>
 License-Expression: MIT

{tskit-1.0.2 → tskit-1.0.3}/lib/tskit/core.c

@@ -121,6 +121,7 @@ tsk_generate_uuid(char *dest, int TSK_UNUSED(flags))
 out:
     return ret;
 }
+
 static const char *
 tsk_strerror_internal(int err)
 {

{tskit-1.0.2 → tskit-1.0.3}/tests/test_drawing.py

@@ -2929,11 +2929,9 @@ class TestDrawKnownSvg(TestDrawSvgBase):
         # set units
         tables.time_units = "generations"
         ts = tables.tree_sequence()
-        svg = ts.draw_svg(y_axis=True, debug_box=draw_plotbox)
+        svg = ts.draw_svg(y_axis=True, title="Y axis test", debug_box=draw_plotbox)
         assert "Time ago (generations)" in svg
-        self.verify_known_svg(
-            svg, "ts_y_axis.svg", overwrite_viz, width=200 * ts.num_trees
-        )
+        self.verify_known_svg(svg, "ts_y_axis.svg", True, width=200 * ts.num_trees)
 
     def test_known_svg_ts_y_axis_regular(self, overwrite_viz, draw_plotbox):
         # This should have gridlines

{tskit-1.0.2 → tskit-1.0.3}/tests/test_metadata.py

@@ -627,6 +627,237 @@ class TestJSONCodec:
         assert ms.decode_row(b"") == {}
 
 
+class TestJSONStructCodec:
+    def test_requires_subschemas(self):
+        with pytest.raises(
+            tskit.MetadataSchemaValidationError,
+            match="requires 'json' and 'struct' schema mappings",
+        ):
+            tskit.MetadataSchema({"codec": "json+struct"})
+
+    def test_disallow_duplicate_keys(self):
+        schema = {
+            "codec": "json+struct",
+            "json": {"type": "object", "properties": {"x": {"type": "number"}}},
+            "struct": {
+                "type": "object",
+                "properties": {"x": {"type": "number", "binaryFormat": "i"}},
+            },
+        }
+        with pytest.raises(
+            tskit.MetadataSchemaValidationError, match="must not share property names"
+        ):
+            tskit.MetadataSchema(schema)
+
+    def test_round_trip_with_struct_and_json(self):
+        schema = {
+            "codec": "json+struct",
+            "json": {
+                "type": "object",
+                "properties": {
+                    "label": {"type": "string"},
+                    "count": {"type": "number"},
+                    "stuff": {"type": "array"},
+                },
+                "required": ["label"],
+            },
+            "struct": {
+                "type": "object",
+                "properties": {
+                    "b0": {"type": "number", "binaryFormat": "i"},
+                    "b1": {"type": "number", "binaryFormat": "i"},
+                    "xyz": {
+                        "type": "array",
+                        "arrayLengthFormat": "H",
+                        "items": {"type": "number", "binaryFormat": "i"},
+                    },
+                },
+            },
+        }
+        ms = tskit.MetadataSchema(schema)
+        for v in [[], [0, 2, 12], [5] * 1000]:
+            row = {
+                "label": "abcdef xyz",
+                "count": 7,
+                "b0": 123,
+                "b1": 0,
+                "stuff": [1, 3, 2, -1.5, "abc", None],
+                "xyz": v,
+                "another_thing": "since JSON is permissive this is allowed",
+            }
+            encoded = ms.validate_and_encode_row(row)
+            out = ms.decode_row(encoded)
+            assert out == row
+
+    def schema_with_binary(self, num_binary_ints):
+        # produces a json+struct schema having num_binary_ints integers
+        # encoded in binary, labeled b0, ... bX
+        schema = {
+            "codec": "json+struct",
+            "json": {
+                "type": "object",
+                "properties": {
+                    "label": {"type": "string"},
+                    "count": {"type": "number"},
+                },
+                "required": ["label"],
+            },
+            "struct": {
+                "type": "object",
+                "properties": {
+                    f"b{j}": {
+                        "type": "integer",
+                        "binaryFormat": "i",
+                    }
+                    for j in range(num_binary_ints)
+                },
+            },
+        }
+        return tskit.MetadataSchema(schema)
+
+    @pytest.mark.parametrize("k", (0, 1, 5, 1001))
+    def test_byte_alignment(self, k):
+        # We want to test whether the binary portion begins byte-aligned.
+        # To verify this, we (somewhat pedantically) let:
+        # X = (bytes to encode the json without any binary)
+        # Y = (bytes to encode the same json with k ints in binary)
+        # and then:
+        # (a) if padding is correct, X should be divisible by 8;
+        # but to make sure that in fact the binary portion starts after X bytes,
+        # we also check that:
+        # (b) Y-X is equal to k * (bytes per int)
+        ms = self.schema_with_binary(k)
+        ms0 = self.schema_with_binary(0)
+        bytes_per_int = len(struct.pack("i", 0))
+        for s in [
+            "",
+            "a",
+            "ab",
+            "abc",
+            "abcd",
+            "abcde",
+            "abcdef",
+            "abcdefg",
+            "abcdefgh",
+            " " * 1000 + "foo" + " " * 1000,
+        ]:
+            row = {"label": s, "count": 7}
+            encoded0 = ms0.validate_and_encode_row(row)
+            row.update({f"b{j}": j for j in range(k)})
+            encoded = ms.validate_and_encode_row(row)
+            out = ms.decode_row(encoded)
+            assert out == row
+            # validate byte alignment
+            assert len(encoded0) % 8 == 0
+            assert len(encoded) - len(encoded0) == k * bytes_per_int
+
+    def test_json_defaults_applied(self):
+        schema = {
+            "codec": "json+struct",
+            "json": {
+                "type": "object",
+                "properties": {"number": {"type": "number", "default": 5}},
+            },
+            "struct": {"type": "object", "properties": {}},
+        }
+        ms = tskit.MetadataSchema(schema)
+        assert ms.decode_row(ms.validate_and_encode_row({})) == {"number": 5}
+        assert ms.decode_row(ms.validate_and_encode_row({"number": 9})) == {"number": 9}
+
+    def test_nested_default_error(self):
+        schema = {
+            "codec": "json+struct",
+            "json": {
+                "type": "object",
+                "properties": {
+                    "obj": {
+                        "type": "object",
+                        "properties": {
+                            "nested_obj_no_default": {
+                                "type": "object",
+                                "properties": {},
+                            },
+                            "nested_obj": {
+                                "type": "object",
+                                "properties": {},
+                                "default": {"foo": "bar"},
+                            },
+                        },
+                    }
+                },
+            },
+            "struct": {"type": "object", "properties": {}},
+        }
+        with pytest.raises(
+            tskit.MetadataSchemaValidationError,
+            match="Defaults can only be specified at the top level for JSON codec",
+        ):
+            tskit.MetadataSchema(schema)
+
+    def test_decode_without_magic_errors(self):
+        ms = tskit.MetadataSchema(
+            {
+                "codec": "json+struct",
+                "json": {"type": "object", "properties": {}},
+                "struct": {"type": "object", "properties": {}},
+            }
+        )
+        with pytest.raises(ValueError, match="missing magic header"):
+            ms.decode_row(b"{}")
+
+    def test_decode_version_mismatch(self):
+        ms = tskit.MetadataSchema(
+            {
+                "codec": "json+struct",
+                "json": {"type": "object", "properties": {}},
+                "struct": {"type": "object", "properties": {}},
+            }
+        )
+        header = metadata.JSONStructCodec._HDR.pack(
+            metadata.JSONStructCodec.MAGIC,
+            metadata.JSONStructCodec.VERSION + 1,
+            len(b"{}"),
+            0,
+        )
+        with pytest.raises(
+            ValueError,
+            match="Unsupported json\\+struct version",
+        ):
+            ms.decode_row(header + b"{}")
+
+    def test_decode_truncated_lengths(self):
+        schema = {
+            "codec": "json+struct",
+            "json": {"type": "object", "properties": {}},
+            "struct": {"type": "object", "properties": {}},
+        }
+        ms = tskit.MetadataSchema(schema)
+        header = metadata.JSONStructCodec._HDR.pack(
+            metadata.JSONStructCodec.MAGIC, metadata.JSONStructCodec.VERSION, 5, 0
+        )
+        with pytest.raises(ValueError, match="declared lengths exceed buffer size"):
+            ms.decode_row(header + b"abc")
+
+        header = metadata.JSONStructCodec._HDR.pack(
+            metadata.JSONStructCodec.MAGIC, metadata.JSONStructCodec.VERSION, 1, 3
+        )
+        with pytest.raises(ValueError, match="declared lengths exceed buffer size"):
+            ms.decode_row(header + b"a")
+
+    def test_missing_struct_property_fails_validation(self):
+        schema = {
+            "codec": "json+struct",
+            "json": {"type": "object", "properties": {}},
+            "struct": {
+                "type": "object",
+                "properties": {"payload": {"type": "integer", "binaryFormat": "i"}},
+            },
+        }
+        ms = tskit.MetadataSchema(schema)
+        with pytest.raises(tskit.MetadataValidationError, match="required property"):
+            ms.validate_and_encode_row({})
+
+
 class TestStructCodec:
     def encode_decode(self, method_name, sub_schema, obj, buffer):
         assert (

{tskit-1.0.2 → tskit-1.0.3}/tests/test_phylo_formats.py

@@ -26,6 +26,7 @@ Tests for phylogenetics export functions, newick, nexus, FASTA etc.
 
 import functools
 import io
+import random
 import textwrap
 
 import dendropy
@@ -334,6 +335,97 @@ class TestNexusIncludeSections:
         )
 
 
+class TestNexusNodeLabels:
+    @tests.cached_example
+    def balanced_tree(self):
+        #   4
+        # ┏━┻┓
+        # ┃  3
+        # ┃ ┏┻┓
+        # 0 1 2
+        return tskit.Tree.generate_balanced(3)
+
+    def test_as_nexus_labels_basic(self):
+        ts = self.balanced_tree().tree_sequence
+        labels = {0: "human", 1: "chimp", 2: "bonobo"}
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=3;
+              TAXLABELS human chimp bonobo;
+            END;
+            BEGIN TREES;
+              TRANSLATE n0 human, n1 chimp, n2 bonobo;
+              TREE t0^1 = [&R] (n0:2,(n1:1,n2:1):1);
+            END;
+            """
+        )
+        assert expected == ts.as_nexus(include_alignments=False, node_labels=labels)
+
+    def test_as_nexus_labels_partial(self):
+        ts = self.balanced_tree().tree_sequence
+        labels = {0: "human", 2: "bonobo"}
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=3;
+              TAXLABELS human n1 bonobo;
+            END;
+            BEGIN TREES;
+              TRANSLATE n0 human, n2 bonobo;
+              TREE t0^1 = [&R] (n0:2,(n1:1,n2:1):1);
+            END;
+            """
+        )
+        assert expected == ts.as_nexus(include_alignments=False, node_labels=labels)
+
+    def test_as_nexus_labels_none(self):
+        ts = self.balanced_tree().tree_sequence
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=3;
+              TAXLABELS n0 n1 n2;
+            END;
+            BEGIN TREES;
+              TREE t0^1 = [&R] (n0:2,(n1:1,n2:1):1);
+            END;
+            """
+        )
+        assert expected == ts.as_nexus(include_alignments=False, node_labels=None)
+
+    @pytest.mark.parametrize("ts", get_example_tree_sequences())
+    def test_parseable(self, ts):
+        def all_samples_are_leaves(ts):
+            internal_nodes = np.unique(ts.edges_parent)
+            is_internal_sample = np.isin(ts.samples(), internal_nodes)
+            return not np.any(is_internal_sample)
+
+        if not all_samples_are_leaves(ts):
+            # TRANSLATE doesn't support translating internal nodes
+            return
+
+        for tree in ts.trees():
+            if not tree.has_single_root:
+                return
+
+        labels = {}
+        samples = ts.samples()
+        k = random.randint(1, len(samples))
+        for node in random.sample(list(samples), k):
+            labels[node] = f"new_node_which_was_{node}"
+
+        nexus = ts.as_nexus(include_alignments=False, node_labels=labels)
+        ds = dendropy.DataSet.get(data=nexus, schema="nexus")
+        tree = ds.tree_lists[0][0]
+        dendropy_labels = [node.taxon.label for node in tree.nodes() if node.taxon]
+        for label in labels.values():
+            assert label.replace("_", " ") in dendropy_labels
+
+
 class TestNewickCodePaths:
     """
     Test that the different code paths we use under the hood lead to

{tskit-1.0.2 → tskit-1.0.3}/tskit/_version.py

@@ -1,4 +1,4 @@
 # Definitive location for the version number.
 # During development, should be x.y.z.devN
 # For beta should be x.y.zbN
-tskit_version = "1.0.2"
+tskit_version = "1.0.3"

{tskit-1.0.2 → tskit-1.0.3}/tskit/drawing.py

@@ -283,6 +283,15 @@ class Timescaling:
         y_scale = self.plot_range / (self.max_time - self.min_time)
         return self.plot_min - (y - self.min_time) * y_scale
 
+    def with_offset(self, y_offset):
+        return Timescaling(
+            max_time=self.max_time,
+            min_time=self.min_time,
+            plot_min=self.plot_min + y_offset,
+            plot_range=self.plot_range,
+            use_log_transform=self.use_log_transform,
+        )
+
 
 class SVGString(str):
     "A string containing an SVG representation"
@@ -1460,7 +1469,7 @@ class SvgTreeSequence(SvgAxisPlot):
                         size=subplot_size, num_skipped=tree.index - last_used_index
                     )
                 )
-        y = self.plotbox.top
+        y_offset = self.plotbox.top
         if title is not None:
             self.add_text_in_group(
                 title,
@@ -1481,14 +1490,16 @@ class SvgTreeSequence(SvgAxisPlot):
             tick_length_upper=self.default_tick_length_site,  # TODO - parameterize
             x_regions=x_regions,
         )
-        y_low = self.tree_plotbox.bottom
-        if y_axis is not None:
+        y_low = self.tree_plotbox.bottom + y_offset
+        if self.y_axis:
             tscales = {s.timescaling for s in subplots if s.timescaling}
             if len(tscales) > 1:
                 raise ValueError(
                     "Can't draw a tree sequence Y axis if trees vary in timescale"
                 )
-            self.timescaling = tscales.pop()
+            # The timescaling of all subplots is used for outer box, but we
+            # need to shift it by the top padding to account for e.g. titles
+            self.timescaling = tscales.pop().with_offset(y_offset)
             y_low = self.timescaling.transform(self.timescaling.min_time)
             if y_ticks is None:
                 used_nodes = edge_and_sample_nodes(ts, breaks[skipbreaks])
@@ -1499,11 +1510,11 @@ class SvgTreeSequence(SvgAxisPlot):
 
         self.draw_y_axis(
             ticks=check_y_ticks(y_ticks),
-            upper=self.tree_plotbox.top,
+            upper=self.tree_plotbox.top + y_offset,
             lower=y_low,
             tick_length_outer=self.default_tick_length,
             gridlines=y_gridlines,
-            side="right" if y_axis == "right" else "left",
+            side="right" if self.y_axis == "right" else "left",
         )
 
         subplot_x = self.plotbox.left
@@ -1513,7 +1524,7 @@ class SvgTreeSequence(SvgAxisPlot):
             svg_subplot = container.add(
                 self.drawing.g(
                     class_=subplot.svg_class,
-                    transform=f"translate({rnd(subplot_x)} {y})",
+                    transform=f"translate({rnd(subplot_x)} {y_offset})",
                 )
             )
             for svg_items in subplot.root_groups.values():

{tskit-1.0.2 → tskit-1.0.3}/tskit/metadata.py

@@ -195,6 +195,149 @@ class NOOPCodec(AbstractMetadataCodec):
         return data
 
 
+class JSONStructCodec(AbstractMetadataCodec):
+    """
+    Pack canonical JSON metadata together with a struct-encoded binary payload.
+    The codec expects a metadata schema with separate ``json`` and ``struct``
+    subschemas and produces a single dict containing the union of the keys from
+    those subschemas after decoding.
+
+    The structure of the encoded metadata is as follows: first, a fixed-size
+    header that contains the number of bytes for both ``json`` and ``struct``
+    portions; next, the json; then a variable number of zeroed padding bytes
+    that brings the length to a multiple of 8 for alignment; and finally
+    the struct-encoded binary portion.
+    """
+
+    MAGIC = b"JBLB"
+    VERSION = 1
+    _HDR = struct.Struct("<4sBQQ")  # magic, version, json_len, blob_len
+
+    @classmethod
+    def is_schema_trivial(cls, schema: Mapping) -> bool:
+        return False
+
+    def __init__(self, schema: Mapping[str, Any]) -> None:
+        json_schema = schema.get("json")
+        struct_schema = schema.get("struct")
+        if not isinstance(json_schema, Mapping) or not isinstance(
+            struct_schema, Mapping
+        ):
+            raise exceptions.MetadataSchemaValidationError(
+                "json+struct requires 'json' and 'struct' schema mappings"
+            )
+
+        json_schema = copy.deepcopy(dict(json_schema, codec="json"))
+        struct_schema = copy.deepcopy(dict(struct_schema, codec="struct"))
+
+        try:
+            json_schema = JSONCodec.modify_schema(json_schema)
+            JSONCodec.schema_validator.check_schema(json_schema)
+        except jsonschema.exceptions.SchemaError as ve:
+            raise exceptions.MetadataSchemaValidationError(str(ve)) from ve
+        try:
+            struct_schema = StructCodec.modify_schema(struct_schema)
+            StructCodecSchemaValidator.check_schema(struct_schema)
+        except jsonschema.exceptions.SchemaError as ve:
+            raise exceptions.MetadataSchemaValidationError(str(ve)) from ve
+
+        self.json_schema = json_schema
+        self.struct_schema = struct_schema
+        json_props = self.json_schema.get("properties", {})
+        struct_props = self.struct_schema.get("properties", {})
+        overlap = set(json_props).intersection(struct_props)
+        if overlap:
+            raise exceptions.MetadataSchemaValidationError(
+                "json and struct schemas must not share property names: "
+                + ", ".join(sorted(overlap))
+            )
+        for name, sub_schema in (
+            ("json", self.json_schema),
+            ("struct", self.struct_schema),
+        ):
+            sub_type = sub_schema.get("type")
+            if sub_type is None:
+                continue
+            if isinstance(sub_type, list):
+                is_object = "object" in sub_type
+            else:
+                is_object = sub_type == "object"
+            if not is_object:
+                raise exceptions.MetadataSchemaValidationError(
+                    f"{name} subschema must describe an object for json+struct codec"
+                )
+
+        self.json_codec = JSONCodec(self.json_schema)
+        self.struct_codec = StructCodec(self.struct_schema)
+        self._struct_keys = set(struct_props.keys())
+        self._validate_json = TSKITMetadataSchemaValidator(self.json_schema).validate
+        self._validate_struct = TSKITMetadataSchemaValidator(self.struct_schema).validate
+
+    def validate_row(self, row: Any) -> None:
+        if not isinstance(row, dict):
+            raise exceptions.MetadataValidationError(
+                "json+struct metadata must be a mapping"
+            )
+        struct_data = {k: v for k, v in row.items() if k in self._struct_keys}
+        json_data = {k: v for k, v in row.items() if k not in self._struct_keys}
+        try:
+            self._validate_json(json_data)
+            self._validate_struct(struct_data)
+        except jsonschema.exceptions.ValidationError as ve:
+            raise exceptions.MetadataValidationError(str(ve)) from ve
+
+    def encode(self, obj: Any) -> bytes:
+        if not isinstance(obj, dict):
+            raise exceptions.MetadataEncodingError(
+                "json+struct metadata must be a mapping"
+            )
+        json_bytes = self.json_codec.encode(
+            {k: v for k, v in obj.items() if k not in self._struct_keys}
+        )
+        blob_bytes = self.struct_codec.encode(
+            {k: v for k, v in obj.items() if k in self._struct_keys}
+        )
+        header = self._HDR.pack(
+            self.MAGIC, self.VERSION, len(json_bytes), len(blob_bytes)
+        )
+        padding_bytes = bytes((8 - ((len(header) + len(json_bytes)) % 8)) % 8)
+        return header + json_bytes + padding_bytes + blob_bytes
+
+    def decode(self, encoded: bytes) -> Any:
+        if len(encoded) >= self._HDR.size and encoded[:4] == self.MAGIC:
+            _, version, jlen, blen = self._HDR.unpack_from(encoded)
+            if version != self.VERSION:
+                raise ValueError("Unsupported json+struct version")
+            start = self._HDR.size
+            # the second mod 8 is here because if start+jlen % 8 is 0 we get 8, not 0
+            padding_length = (8 - ((start + jlen) % 8)) % 8
+            if (
+                jlen > len(encoded) - start
+                or blen > len(encoded) - start - jlen - padding_length
+            ):
+                raise ValueError(
+                    "Invalid json+struct payload: declared lengths exceed buffer size"
+                )
+            json_bytes = encoded[start : start + jlen]
+            blob_bytes = encoded[
+                start + jlen + padding_length : start + jlen + padding_length + blen
+            ]
+            json_data = self.json_codec.decode(json_bytes)
+            struct_data = self.struct_codec.decode(blob_bytes)
+            overlap = set(json_data).intersection(struct_data)
+            if overlap:
+                raise ValueError(
+                    "json+struct decoded duplicate keys: " + ", ".join(sorted(overlap))
+                )
+            combined = dict(json_data)
+            combined.update(struct_data)
+            return combined
+        raise ValueError("Invalid json+struct payload: missing magic header")
+
+
+register_metadata_codec(JSONStructCodec, "json+struct")
+
+
 def binary_format_validator(validator, types, instance, schema):
     # We're hooking into jsonschemas validation code here, which works by creating
     # generators of exceptions, hence the yielding
@@ -805,7 +948,11 @@ class MetadataSchema:
             self._schema = codec_cls.modify_schema(schema)
             self.codec_instance = codec_cls(self._schema)
             self._string = tskit.canonical_json(self._schema)
-            self._validate_row = TSKITMetadataSchemaValidator(self._schema).validate
+            self._validate_row = getattr(
+                self.codec_instance,
+                "validate_row",
+                TSKITMetadataSchemaValidator(self._schema).validate,
+            )
             self._bypass_validation = codec_cls.is_schema_trivial(schema)
             self.encode_row = self.codec_instance.encode
             self.decode_row = self.codec_instance.decode

{tskit-1.0.2 → tskit-1.0.3}/tskit/tables.py

@@ -3455,6 +3455,24 @@ class TableCollection(metadata.MetadataProvider):
 
     @classmethod
     def load(cls, file_or_path, *, skip_tables=False, skip_reference_sequence=False):
+        """
+        Load a :class:`TableCollection` from a file or path, saved in the format defined
+        by :meth:`.dump`. Although the file must be in the correct format, unlike
+        :func:`tskit.load` further validation is not performed to check that the tables
+        meet the :ref:`sec_valid_tree_sequence_requirements`. For instance,
+        unsorted tables can be loaded without error and then sorted before
+        :meth:`creating <TableCollection.tree_sequence>` a :class:`TreeSequence`.
+
+        :param file_or_path: The file object or path from which to load the
+            TableCollection.
+        :param bool skip_tables: If True, no tables are read from the file and
+            only the top-level information is populated in the returned TableCollection.
+        :param bool skip_reference_sequence: If True, tables are read
+            without loading any reference sequence.
+        :return: A TableCollection instance
+        :rtype: TableCollection
+        :raises: **tskit.FileFormatError** -- If the file is not in a valid format.
+        """
         file, local_file = util.convert_file_like_to_open_file(file_or_path, "rb")
         ll_tc = _tskit.TableCollection()
         try:
@@ -3474,7 +3492,8 @@ class TableCollection(metadata.MetadataProvider):
         """
         Writes the table collection to the specified path or file object.
 
-        :param str file_or_path: The file object or path to write the TreeSequence to.
+        :param str file_or_path: The file object or path to which to write this
+            TableCollection.
         """
         file, local_file = util.convert_file_like_to_open_file(file_or_path, "wb")
         try:

{tskit-1.0.2 → tskit-1.0.3}/tskit/text_formats.py

@@ -120,6 +120,7 @@ def write_nexus(
     include_alignments,
     reference_sequence,
     missing_data_character,
+    node_labels,
     isolated_as_missing=None,
 ):
     # See TreeSequence.write_nexus for documentation on parameters.
@@ -134,7 +135,13 @@ def write_nexus(
     print("#NEXUS", file=out)
     print("BEGIN TAXA;", file=out)
     print("", f"DIMENSIONS NTAX={ts.num_samples};", sep=indent, file=out)
-    taxlabels = " ".join(f"n{u}" for u in ts.samples())
+
+    if node_labels is not None:
+        taxlabels = " ".join(
+            node_labels[u] if u in node_labels else f"n{u}" for u in ts.samples()
+        )
+    else:
+        taxlabels = " ".join(f"n{u}" for u in ts.samples())
     print("", f"TAXLABELS {taxlabels};", sep=indent, file=out)
     print("END;", file=out)
 
@@ -166,6 +173,11 @@ def write_nexus(
     include_trees = True if include_trees is None else include_trees
     if include_trees:
         print("BEGIN TREES;", file=out)
+
+        if node_labels is not None:
+            translations = ", ".join(f"n{u} {name}" for u, name in node_labels.items())
+            print(f"  TRANSLATE {translations};", file=out)
+
         for tree in ts.trees():
             start_interval = "{0:.{1}f}".format(tree.interval.left, pos_precision)
             end_interval = "{0:.{1}f}".format(tree.interval.right, pos_precision)

{tskit-1.0.2 → tskit-1.0.3}/tskit/trees.py

@@ -6797,6 +6797,7 @@ class TreeSequence:
         reference_sequence=None,
         missing_data_character=None,
         isolated_as_missing=None,
+        node_labels=None,
     ):
         """
         Returns a `nexus encoding <https://en.wikipedia.org/wiki/Nexus_file>`_
@@ -6896,6 +6897,10 @@ class TreeSequence:
         :param str missing_data_character: As for the :meth:`.alignments` method,
             but defaults to "?".
         :param bool isolated_as_missing: As for the :meth:`.alignments` method.
+        :param node_labels: A map of type `{node_id: name}`.  Samples present
+            in the map will have the given name instead of `n{node_id}`.  Note
+            that the names must not have whitespace (spaces should be replaced
+            by underscores) or puncuation in them.
         :return: A nexus representation of this :class:`TreeSequence`
         :rtype: str
         """
@@ -6908,6 +6913,7 @@ class TreeSequence:
             reference_sequence=reference_sequence,
             missing_data_character=missing_data_character,
             isolated_as_missing=isolated_as_missing,
+            node_labels=node_labels,
         )
 
     def as_nexus(self, **kwargs):
@@ -10930,12 +10936,106 @@ class TreeSequence:
     def ld_matrix(
         self,
         sample_sets=None,
-        sites=None,
-        positions=None,
         mode="site",
         stat="r2",
+        sites=None,
+        positions=None,
         indexes=None,
     ):
+        r"""
+
+        Returns a matrix of the specified two-locus statistic (default
+        :math:`r^2`) computed from sample allelic states or branch lengths.
+        The resulting linkage disequilibrium (LD) matrix represents either the
+        two-locus statistic as computed between all pairs of specified
+        ``sites`` (``"site"`` mode, producing a
+        ``len(sites)``-by-``len(sites)`` sized matrix), or as computed from the
+        branch structures at marginal trees between pairs of trees at all
+        specified ``positions`` (``"branch"`` mode, producing a
+        ``len(positions)``-by-``len(positions)`` sized matrix).
+
+        The sites considered for ``"site"`` mode defaults to all sites (which may
+        result in a very large matrix!), but can be restricted using
+        the ``sites`` argument. Sites must be passed as a list of lists,
+        specifying the ``[row_sites, col_sites]``, resulting in a
+        rectangular matrix, or by specifying a single list of ``[sites]``, in
+        which a square matrix will be produced (see
+        :ref:`sec_stats_two_locus_site` for examples). Here, ``sites``,
+        ``row_sites``, and ``col_sites`` are each lists of site indexes.
+
+        Similarly, in the ``"branch"`` mode, the ``positions`` argument specifies
+        genomic coordinates at which the expectation for the two-locus statistic
+        is computed, given the local tree structure.
+        (See :ref:`sec_stats_two_locus_branch` for explanation of in what sense
+        this is an expectation.) This defaults to computing
+        the LD for each pair of distinct trees (this is equivalent to passing in
+        the leftmost coordinates of each tree's span, since intervals are closed on
+        the left and open on the right). Similar to the site mode, a nested list
+        of row and column positions can be specified separately (resulting in a
+        rectangular matrix) or a single list of a specified positions results
+        in a square matrix (see :ref:`sec_stats_two_locus_branch` for
+        examples). Like ``sites``, the ``positions`` must be specified as a list
+        of lists.
+
+        Some LD statistics are defined for both within a single set of samples
+        and for two sample sets. If the ``indexes`` argument is specified, then
+        ``indexes`` specifies the indexes of the sample sets in the
+        ``sample_sets`` list between which to compute LD. For instance, this
+        results in a 3D array whose ``[k,:,:]``-th slice contains LD values
+        between ``sample_sets[i]`` and ``sample_sets[j]``, where ``(i, j)`` is
+        the ``k``-th element of ``indexes``.
+
+        For more on how the ``indexes`` and ``sample_sets`` interact with the
+        output dimensions, see the :ref:`sec_stats_two_locus_sample_sets`
+        section. Statistics are defined in the
+        :ref:`sec_stats_two_locus_summary_functions_two_way` section.
+
+        **Available Stats** (use ``Stat Name`` in the ``stat`` keyword
+        argument). Statistics marked as "multi sample set" allow
+        (but do not require) computation from two sample sets
+        via the ``indexes`` argument.
+
+        ======================= ========== ================ ==============
+        Stat                     Polarised Multi Sample Set Stat Name
+        ======================= ========== ================ ==============
+        :math:`r^2`              n          y               "r2"
+        :math:`r`                y          n               "r"
+        :math:`D^2`              n          y               "D2"
+        :math:`D`                y          n               "D"
+        :math:`D'`               y          n               "D_prime"
+        :math:`D_z`              n          n               "Dz"
+        :math:`\pi_2`            n          n               "pi2"
+        :math:`\widehat{D^2}`    n          y               "D2_unbiased"
+        :math:`\widehat{D_z}`    n          n               "Dz_unbiased"
+        :math:`\widehat{\pi_2}`  n          n               "pi2_unbiased"
+        ======================= ========== ================ ==============
+
+        :param list sample_sets: A list, or a list of lists of sample node IDs,
+            specifying the groups of nodes to compute the statistic with. Defaults
+            to all samples.
+        :param str mode: A string giving the "type" of the statistic to be
+            computed. Defaults to "site", can be "site" or "branch".
+        :param str stat: A string giving the selected two-locus statistic to
+            compute. Defaults to "r2".
+        :param list sites: A list of lists of sites over which to compute an
+            LD matrix. Can be specified as a list of lists to control the row
+            and column sites. Only available in "site" mode. Specify as
+            ``[row_sites, col_sites]`` or ``[all_sites]``.
+            Defaults to all sites.
+        :param list positions: A list of lists of genomic positions where
+            expected LD is computed based on tree topologies and branch
+            lengths. Only applicable in "branch" mode. Specify as a list of
+            two lists to control the row and column positions, as
+            ``[row_positions, col_positions]``, or ``[all_positions]``.
+            Defaults to the leftmost coordinates of all trees and computes
+            LD between all pairs of trees.
+        :param list indexes: A list of 2-tuples or a single 2-tuple, specifying
+            the indexes of two sample sets over which to compute a two-way LD
+            statistic. Only :math:`r^2`, :math:`D^2`, and :math:`\widehat{D^2}`
+            are implemented for two-way statistics.
+        :return: A 2D or 3D array of LD matrices.
+        :rtype: numpy.ndarray
+        """
         one_way_stats = {
             "D": self._ll_tree_sequence.D_matrix,
             "D2": self._ll_tree_sequence.D2_matrix,

{tskit-1.0.2 → tskit-1.0.3/tskit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tskit
-Version: 1.0.2
+Version: 1.0.3
 Summary: The tree sequence toolkit.
 Author-email: Tskit Developers <admin@tskit.dev>
 License-Expression: MIT