def test_slicing_combos(self):
"""different mixtures of slicing should work"""
darr = DictArrayTemplate(list(DNA), list(DNA)).wrap([
[0.7, 0.1, 0.2, 0.3],
[0.1, 0.7, 0.1, 0.3],
[0.3, 0.2, 0.6, 0.3],
[0.4, 0.1, 0.1, 0.7],
])
got = darr["C":"G", "C":"G"]
assert_allclose(got.array, numpy.array([[0.7, 0.1], [0.2, 0.6]]))
got = darr[[1, 2], [1, 2]]
assert_allclose(got.array, numpy.array([[0.7, 0.1], [0.2, 0.6]]))
got = darr[[2, 3], "C"]
assert_allclose(got.array, numpy.array([0.2, 0.1]))
got = darr["C", [2, 3]]
assert_allclose(got.array, numpy.array([0.1, 0.3]))
got = darr[[1, 2], "T":"A"]
assert_allclose(got.array, numpy.array([[0.1, 0.7], [0.3, 0.2]]))
got = darr["T":"A", [1, 2]]
assert_allclose(got.array, numpy.array([[0.1, 0.2], [0.7, 0.1]]))
# make sure we cope with keys that are int's
nums = list(range(1, 5))
darr = DictArrayTemplate(nums, nums).wrap([
[0.7, 0.1, 0.2, 0.3],
[0.1, 0.7, 0.1, 0.3],
[0.3, 0.2, 0.6, 0.3],
[0.4, 0.1, 0.1, 0.7],
])
got = darr[[1, 2], [1, 2]]
assert_allclose(got.array, numpy.array([[0.7, 0.1], [0.2, 0.6]]))
def test_convert2DDict(self):
"""convert2DDict produces valid template input"""
data = dict(a=dict(b=4, c=5))
vals, row_keys, col_keys = convert2DDict(data)
self.assertEqual(set(row_keys), set(["a"]))
b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
self.assertEqual(b.array.tolist(), [[4, 5]])
# row keys, then column
self.assertEqual(b.template.names, [["a"], ["b", "c"]])
data = {
"a": {"a": 0, "b": 1, "e": 0},
"b": {"a": 1, "b": 0, "e": 4},
"e": {"a": 0, "b": 4, "e": 0},
}
vals, row_keys, col_keys = convert2DDict(data)
b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
got = b.to_dict()
self.assertEqual(got, data)
self.assertEqual(b.template.names, [["a", "b", "e"], ["a", "b", "e"]])
data = dict(a=dict(b=4, c=5))
vals, row_keys, col_keys = convert2DDict(data, make_symmetric=True)
self.assertEqual(row_keys, col_keys)
self.assertEqual(vals, [[0, 4, 5], [4, 0, 0], [5, 0, 0]])
def to_dictarray(self):
"""construct fully enumerated dictarray
Returns
-------
DictArray with dtype of int
Notes
-----
Unobserved combinations have zeros. Result can can be indexed as if it was a numpy array using key values
"""
from itertools import product
from cogent3.util.dict_array import DictArrayTemplate
key = next(iter(self))
try:
ndim = 1 if isinstance(key, str) else len(key)
except TypeError:
ndim = 1
if ndim == 1:
names = sorted(self)
vals = [self[n] for n in names]
darr = DictArrayTemplate(names).wrap(vals, dtype=int)
return darr
categories = [sorted(set(labels)) for labels in zip(*self)]
shape = tuple(len(c) for c in categories)
darr = DictArrayTemplate(*categories).wrap(numpy.zeros(shape, dtype=int))
for comb in product(*categories):
indices = [[categories[i].index(c)] for i, c in enumerate(comb)]
darr.array[tuple(indices)] = self[comb]
return darr
def test_to_dict_nested(self):
"""DictArray.to_dict() should convert nested DictArray instances to
dict's too."""
a = numpy.identity(3, int)
b = DictArrayTemplate("abc", "ABC")
b = b.wrap(a)
self.assertEqual(b.array.tolist(), [[1, 0, 0], [0, 1, 0], [0, 0, 1]])
c = DictArrayTemplate("de", "DE").wrap([[b, b], [b, b]])
self.assertTrue(isinstance(c.to_dict()["d"], dict))
def test_convert_for_dictarray(self):
"""successfully delegates when constructed from a DictArray"""
a = numpy.identity(3, int)
b = DictArrayTemplate("abc", "ABC").wrap(a)
vals, row_keys, col_keys = convert_for_dictarray(b)
got = DictArrayTemplate(row_keys, col_keys).wrap(vals)
self.assertEqual(got.array.tolist(), b.array.tolist())
# the wrap method creates a new array
self.assertIsNot(got.array, b.array)
def test_convert_series(self):
"""convert_series produces valid template input"""
vals, row_keys, col_keys = convert_series([[4], [5]], ["A", "B"], ["a"])
b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
self.assertEqual(b.array.tolist(), [[4], [5]])
data = [[245, 599]]
vals, row_keys, col_keys = convert_series(data)
b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
self.assertEqual(b.array.tolist(), data)
vals, row_keys, col_keys = convert_series(data[0])
b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
self.assertEqual(b.array.tolist(), data[0])
def test_convert_for_dictarray(self):
"""convert_for_dictarray correctly delegates"""
b = DictArrayTemplate("abc", "ABC").wrap(self.a)
data_types = (
[[245, 599]],
dict(a=dict(b=4, c=5)),
{("a", "b"): 4, ("a", "c"): 5},
dict(a=0, b=35, c=45),
b,
)
for data in data_types:
vals, row_keys, col_keys = convert_for_dictarray(data)
_ = DictArrayTemplate(row_keys, col_keys).wrap(vals)
def test_convert_1D_dict(self):
"""convert_1D_dict produces valid template input"""
data = dict(a=0, b=35, c=45)
vals, keys = convert_1D_dict(data)
b = DictArrayTemplate(keys)
b = b.wrap(vals)
self.assertEqual(b.array.tolist(), [0, 35, 45])
def _get_motif_probs_by_node_tr(self, edges=None, bin=None, locus=None):
"""returns motif probs by node for time-reversible models"""
mprob_rules = [
r for r in self.get_param_rules() if "mprob" in r["par_name"]
]
if len(mprob_rules) > 1 or self.model.mprob_model == "monomers":
raise NotImplementedError
mprobs = self.get_motif_probs()
if len(mprobs) != len(self.motifs):
# a Muse and Gaut model
expanded = numpy.zeros(len(self.motifs), dtype=float)
for i, motif in enumerate(self.motifs):
val = 1.0
for b in motif:
val *= mprobs[b]
expanded[i] = val
mprobs = expanded / expanded.sum()
else:
mprobs = [mprobs[m] for m in self.motifs]
edges = []
values = []
for e in self.tree.postorder():
edges.append(e.name)
values.append(mprobs)
return DictArrayTemplate(edges, self.motifs).wrap(values)
def test_get_repr_html(self):
"""exercising method used by parent classes for nice Jupyter display"""
data = [[3, 7], [2, 8], [5, 5]]
darr = DictArrayTemplate(list("ABC"), list("ab")).wrap(data)
got = darr._repr_html_()
self.assertIsInstance(got, str)
self.assertTrue(len(got), 100)
def __init__(self,
default=None,
name=None,
dimensions=None,
dimension=None,
size=None,
**kw):
assert name
if size is not None:
pass
elif default is not None:
size = len(default)
elif dimension is not None:
size = len(dimension[1])
self.size = size
if dimension is not None:
self.internal_dimension = dimension
(dim_name, dim_cats) = dimension
self.bin_names = dim_cats
self.array_template = DictArrayTemplate(dim_cats)
self.internal_dimensions = (dim_name, )
if default is None:
default = self._make_default_value()
elif self.array_template is not None:
default = self.array_template.unwrap(default)
else:
default = numpy.asarray(default)
_InputDefn.__init__(self,
name=name,
default=default,
dimensions=dimensions,
**kw)
self.check_value_is_valid(default, True)
def test_key_levels(self):
"""DictArray both levels have keys."""
b = DictArrayTemplate("abc", "ABC").wrap(self.a)
self.assertEqual(b.keys(), ["a", "b", "c"])
self.assertEqual(b["a"].keys(), ["A", "B", "C"])
self.assertEqual(list(b["a"]), [1, 0, 0])
self.assertEqual(sum(b["a"]), 1)
def get_rate_matrix_for_edge(self, name, calibrated=True, **kw):
"""returns the rate matrix (Q) for the named edge
Parameters
----------
name : str
name of the edge
calibrated : bool
If True, the rate matrix is scaled such that
``sum(pi_i * Qii) == 1``. If False, the calibrated matrix is
multiplied by the length parameter (and the rate parameter for a
bin if it is a rate heterogeneity model).
Notes
-----
If ``calibrated=False``, ``expm(Q)`` will give the same result as
``self.get_psub_for_edge(name)``
"""
# todo handle case of multiple loci
try:
array = self.get_param_value("Q", edge=name, **kw)
array = array.copy()
if not calibrated:
length = self.get_param_value("length", edge=name, **kw)
array *= length
except KeyError as err:
if err[0] == "Q" and name != "Q":
raise RuntimeError("rate matrix not known by this model")
else:
raise
return DictArrayTemplate(self._motifs, self._motifs).wrap(array)
def get_bin_probs(self, locus=None):
hmm = self.get_param_value("bindex", locus=locus)
lhs = [
self.get_param_value("lh", locus=locus, bin=bin) for bin in self.bin_names
]
array = hmm.get_posterior_probs(*lhs)
return DictArrayTemplate(self.bin_names, array.shape[1]).wrap(array)
def get_psub_for_edge(self, name, **kw):
"""returns the substitution probability matrix for the named edge"""
try:
# For PartialyDiscretePsubsDefn
array = self.get_param_value("dpsubs", edge=name, **kw)
except KeyError:
array = self.get_param_value("psubs", edge=name, **kw)
return DictArrayTemplate(self._motifs, self._motifs).wrap(array)
def test_str_labels(self):
"""DictArray with str labels or numpy U dtype"""
b = DictArrayTemplate(["Ab", "Bb", "Cb"],
["Db", "Eb", "Fb"]).wrap(self.a)
self.assertEqual(b.keys(), ["Ab", "Bb", "Cb"])
self.assertEqual(b[0].keys(), ["Db", "Eb", "Fb"])
self.assertEqual(b["Ab", "Eb"], 0)
self.assertEqual(b["Bb", "Eb"], 1)
b = DictArrayTemplate(
numpy.array(["Ab", "Bb", "Cb"], dtype="U"),
numpy.array(["Db", "Eb", "Fb"], dtype="U"),
).wrap(self.a)
self.assertEqual(b.keys(), ["Ab", "Bb", "Cb"])
self.assertEqual(b[0].keys(), ["Db", "Eb", "Fb"])
self.assertEqual(b["Ab", "Eb"], 0)
self.assertEqual(b["Bb", "Eb"], 1)
def test_to_dict(self):
"""DictArray should convert 1D / 2D arrays with/without named row"""
# 1D data, only 1D keys provided
data = [0, 35, 45]
keys = "a", "b", "c"
darr = DictArrayTemplate(keys).wrap(data)
self.assertEqual(darr.to_dict(), dict(zip(keys, data)))
# 2D data, 2D keys, both string, provided
data = [[0, 35, 45]]
darr = DictArrayTemplate(["0"], keys).wrap(data)
darr.to_dict()
self.assertEqual(darr.to_dict(), {"0": {"a": 0, "b": 35, "c": 45}})
# 2D data, 2D keys, one int, one string, provided
darr = DictArrayTemplate([1], keys).wrap(data)
self.assertEqual(darr.to_dict(), {1: {"a": 0, "b": 35, "c": 45}})
darr = DictArrayTemplate([0], keys).wrap(data)
self.assertEqual(darr.to_dict(), {0: {"a": 0, "b": 35, "c": 45}})
def test_interpret_index(self):
"""correctly handles just explicitly defined indices"""
n = ["ab", "dna", "rna"]
a1D = DictArrayTemplate(n)
got = a1D.interpret_index(["ab", "rna"])
self.assertEqual(got[0], ([0, 2], ))
got = a1D.interpret_index([0, 2])
self.assertEqual(got[0], ([0, 2], ))
def test_category_counts_from_non_int_arrays(self):
"""handles object and float numpy array, fails if float"""
a = numpy.array([[31, 36], [58, 138]], dtype=object)
darr = DictArrayTemplate(["syn", "nsyn"], ["Ts", "Tv"]).wrap(a)
got = CategoryCounts(darr)
assert_allclose(got.observed.array.tolist(), a.tolist())
for dtype in (object, float):
with self.assertRaises(TypeError):
a = numpy.array([[31.3, 36], [58, 138]], dtype=dtype)
darr = DictArrayTemplate(["syn", "nsyn"], ["Ts", "Tv"]).wrap(a)
_ = CategoryCounts(darr)
# negative values disallowed
with self.assertRaises(ValueError):
a = numpy.array([[31, -36], [58, 138]], dtype=int)
darr = DictArrayTemplate(["syn", "nsyn"], ["Ts", "Tv"]).wrap(a)
_ = CategoryCounts(darr)
def test_get_repr_html(self):
"""exercising method used by parent classes for nice Jupyter display"""
data = [[3, 7], [2, 8], [5, 5]]
darr = DictArrayTemplate(list("ABC"), list("ab")).wrap(data)
got = darr._repr_html_()
self.assertIsInstance(got, str)
self.assertTrue(len(got), 100)
# case where 1D array
a = [4, 6, 4, 2]
darr = DictArrayTemplate(["A", "C", "G", "T"]).wrap(a)
got = darr._repr_html_()
self.assertTrue('class="index"' not in got)
# case of 3D array
d3 = numpy.arange(8).reshape((2, 2, 2))
darr = DictArrayTemplate(2, 2, 2).wrap(d3)
got = darr._repr_html_()
self.assertIn("3 dimensional", got)
def get_all_rate_matrices(self, calibrated=True):
"""returns all rate matrices (Q) as a dict, keyed by scope
Parameters
----------
calibrated : bool
scales the rate matrix by branch length for each edge. If a rate
heterogeneity model, then the matrix is further scaled by rate
for a bin
Returns
-------
If a single rate matrix, the key is an empty tuple
"""
defn = self.defn_for["Q"]
rate_het = self.defn_for.get("rate", False)
if rate_het:
bin_index = rate_het.valid_dimensions.index("bin")
bin_names = [k[bin_index] for k in rate_het.index]
bin_names = {n: i for i, n in enumerate(bin_names)}
bin_index = defn.valid_dimensions.index("bin")
else:
bin_names = None
bin_index = None
used_dims = defn.used_dimensions()
edge_index = defn.valid_dimensions.index("edge")
indices = {defn.valid_dimensions.index(k) for k in used_dims}
if not calibrated:
indices.add(edge_index)
if not calibrated and rate_het:
indices.add(bin_index)
indices = list(sorted(indices))
result = {}
darr_template = DictArrayTemplate(self._motifs, self._motifs)
for scope, index in defn.index.items():
q = defn.values[index] # this gives the appropriate Q
# from scope we extract only the relevant dimensions
key = tuple(numpy.take(scope, indices))
q = q.copy()
if not calibrated:
length = self.get_param_value("length", edge=scope[edge_index])
if rate_het:
bdex = bin_names[scope[bin_index]]
rate = rate_het.values[bdex]
length *= rate
q *= length
result[key] = darr_template.wrap(q)
if not indices and calibrated:
break # single rate matrix
return result
def test_to_normalized(self):
"""computes frequencies across correct dimension"""
data = [[3, 7], [2, 8], [5, 5]]
darr = DictArrayTemplate(list("ABC"), list("ab")).wrap(data)
row_normal = darr.to_normalized(by_row=True)
assert_allclose(row_normal.array, [[0.3, 0.7], [0.2, 0.8], [0.5, 0.5]])
col_normal = darr.to_normalized(by_column=True)
assert_allclose(col_normal.array, [[0.3, 7 / 20], [0.2, 8 / 20], [0.5, 5 / 20]])
# trying to do both raises AssertionError
with self.assertRaises(AssertionError):
darr.to_normalized(by_row=True, by_column=True)
def test_numpy_ops(self):
"""DictArray should work properly in numpy operations."""
darr = DictArrayTemplate(list(DNA), list(DNA)).wrap([
[0.7, 0.1, 0.1, 0.1],
[0.1, 0.7, 0.1, 0.1],
[0.1, 0.1, 0.7, 0.1],
[0.1, 0.1, 0.1, 0.7],
])
mprobs = numpy.array([0.25, 0.25, 0.25, 0.25])
assert_allclose(mprobs.dot(darr), [0.25, 0.25, 0.25, 0.25])
assert_allclose(numpy.dot(mprobs, darr), [0.25, 0.25, 0.25, 0.25])
def test_get_logo_missing(self):
"""copes with positions with no values"""
data = [
[0.1, 0.3, 0.5, 0.1],
[0.05, 0.8, 0.05, 0.1],
[0, 0, 0, 0],
[0.7, 0.1, 0.1, 0.1],
[0.6, 0.15, 0.05, 0.2],
]
data = DictArrayTemplate(5, "ACGT").wrap(data)
d = get_logo(data)
def test_get_logo(self):
"""returns Drawable"""
data = [
[0.1, 0.3, 0.5, 0.1],
[0.25, 0.25, 0.25, 0.25],
[0.05, 0.8, 0.05, 0.1],
[0.7, 0.1, 0.1, 0.1],
[0.6, 0.15, 0.05, 0.2],
]
data = DictArrayTemplate(5, "ACGT").wrap(data)
d = get_logo(data)
def test_getitem(self):
"""correctly slices"""
darr = DictArrayTemplate(list(DNA), list(DNA)).wrap([
[0.7, 0.1, 0.1, 0.1],
[0.1, 0.7, 0.1, 0.1],
[0.1, 0.1, 0.7, 0.1],
[0.1, 0.1, 0.1, 0.7],
])
r = darr[:, "A":"G"]
assert_allclose(r.to_array(), [[0.1], [0.1], [0.7], [0.1]])
r = darr[2:, "A":"G"]
assert_allclose(r.to_array(), [[0.7], [0.1]])
def test_direct_construction(self):
"""directly construct a dict array"""
b = DictArrayTemplate("abc", "ABC").wrap(self.a)
data_types = (
[[245, 599]],
dict(a=dict(b=4, c=5)),
{("a", "b"): 4, ("a", "c"): 5},
dict(a=0, b=35, c=45),
b,
)
for data in data_types:
g = DictArray(data)
def test_write(self):
"""exercising write method"""
data = [[3, 7], [2, 8], [5, 5]]
darr = DictArrayTemplate(list("ABC"), list("ab")).wrap(data)
with TemporaryDirectory(dir=".") as dirname:
outpath = os.path.join(dirname, "delme.tsv")
darr.write(outpath)
with open(outpath) as infile:
contents = [l.strip().split() for l in infile]
header = contents.pop(0)
self.assertEqual(header, ["dim-1", "dim-2", "value"])
got = {(k1, k2): int(v) for k1, k2, v in contents}
self.assertEqual(got, darr.to_dict(flatten=True))
def reconstruct_ancestral_seqs(self, locus=None):
"""computes the conditional probabilities of each state for each node
in the tree.
Parameters
----------
locus
a named locus
Returns
-------
{node_name: DictArray, ...}
Notes
-----
Alignment columns are rows in the DictArray.
"""
result = {}
array_template = None
for restricted_edge in self._tree.get_edge_vector():
if restricted_edge.istip():
continue
try:
r = []
for motif in range(len(self._motifs)):
self.set_param_rule(
"fixed_motif",
value=motif,
edge=restricted_edge.name,
locus=locus,
is_constant=True,
)
likelihoods = self.get_full_length_likelihoods(locus=locus)
r.append(likelihoods)
if array_template is None:
array_template = DictArrayTemplate(
likelihoods.shape[0], self._motifs
)
finally:
self.set_param_rule(
"fixed_motif",
value=-1,
edge=restricted_edge.name,
locus=locus,
is_constant=True,
)
# dict of site x motif arrays
result[restricted_edge.name] = array_template.wrap(
numpy.transpose(numpy.asarray(r))
)
return result
def get_motif_probs_by_node(self, edges=None, bin=None, locus=None):
from cogent3.evolve.substitution_model import TimeReversible
if isinstance(self.model, TimeReversible):
return self._get_motif_probs_by_node_tr(edges=edges, bin=bin, locus=locus)
kw = dict(bin=bin, locus=locus)
mprobs = self.get_param_value("mprobs", **kw)
mprobs = self._model.calc_word_probs(mprobs)
result = self._nodeMotifProbs(self._tree, mprobs, kw)
if edges is None:
edges = [name for (name, m) in result]
result = dict(result)
values = [result[name] for name in edges]
return DictArrayTemplate(edges, self._mprob_motifs).wrap(values)