def test_import(self, tree, tmpdir):
"""Test import of frequencies JSON that was exported from a frequencies instance.
"""
start_date = 2015.5
end_date = 2018.5
kde_frequencies = TreeKdeFrequencies(start_date=start_date,
end_date=end_date)
frequencies = kde_frequencies.estimate(tree)
frequencies_json = kde_frequencies.to_json()
# Try to dump exported JSON to disk.
tmp_fh = tmpdir.mkdir("json").join("frequencies.json")
fh = tmp_fh.open(mode="w")
json.dump(frequencies_json, fh)
fh.close()
assert tmp_fh.check()
# Import frequencies from existing tree and JSON.
fh = tmp_fh.open()
new_frequencies_json = json.load(fh)
fh.close()
new_kde_frequencies = TreeKdeFrequencies.from_json(
new_frequencies_json)
assert np.array_equal(kde_frequencies.pivots,
new_kde_frequencies.pivots)
# Get the first non-root key (root clade is number 0) and should be first in the sorted list of keys.
key = sorted(kde_frequencies.frequencies.keys())[1]
assert np.array_equal(kde_frequencies.frequencies[key],
new_kde_frequencies.frequencies[key])
def test_import_without_frequencies(self):
"""Test import of frequencies JSON that was exported from a frequencies instance without frequency values.
"""
kde_frequencies = TreeKdeFrequencies()
frequencies_json = kde_frequencies.to_json()
# Import frequencies from existing tree and JSON.
new_kde_frequencies = TreeKdeFrequencies.from_json(frequencies_json)
assert kde_frequencies.pivot_frequency == new_kde_frequencies.pivot_frequency
assert not hasattr(new_kde_frequencies, "frequencies")
def test_import_without_frequencies(self):
"""Test import of frequencies JSON that was exported from a frequencies instance without frequency values.
"""
kde_frequencies = TreeKdeFrequencies()
frequencies_json = kde_frequencies.to_json()
# Import frequencies from existing tree and JSON.
new_kde_frequencies = TreeKdeFrequencies.from_json(frequencies_json)
assert kde_frequencies.pivot_frequency == new_kde_frequencies.pivot_frequency
assert not hasattr(new_kde_frequencies, "frequencies")
def test_import(self, tree, tmpdir):
"""Test import of frequencies JSON that was exported from a frequencies instance.
"""
start_date = 2015.5
end_date = 2018.5
kde_frequencies = TreeKdeFrequencies(
start_date=start_date,
end_date=end_date
)
frequencies = kde_frequencies.estimate(tree)
frequencies_json = kde_frequencies.to_json()
# Try to dump exported JSON to disk.
tmp_fh = tmpdir.mkdir("json").join("frequencies.json")
fh = tmp_fh.open(mode="w")
json.dump(frequencies_json, fh)
fh.close()
assert tmp_fh.check()
# Import frequencies from existing tree and JSON.
fh = tmp_fh.open()
new_frequencies_json = json.load(fh)
fh.close()
new_kde_frequencies = TreeKdeFrequencies.from_json(new_frequencies_json)
assert np.array_equal(
kde_frequencies.pivots,
new_kde_frequencies.pivots
)
# Get the first non-root key (root clade is number 0) and should be first in the sorted list of keys.
key = sorted(kde_frequencies.frequencies.keys())[1]
assert np.array_equal(
kde_frequencies.frequencies[key],
new_kde_frequencies.frequencies[key]
)
parser.add_argument("--frequency-method", required=True, choices=["kde", "diffusion"], help="method used to estimate frequencies")
parser.add_argument("--clades", help="JSON of clade annotations for nodes in the given tree")
parser.add_argument("--delta-pivots", type=int, default=1, help="number of frequency pivots to look back in time for change in frequency calculation")
parser.add_argument("--output", required=True, help="JSON of delta frequency annotations for nodes in the given tree")
args = parser.parse_args()
# Load the tree.
tree = Bio.Phylo.read(args.tree, "newick")
# Load frequencies.
with open(args.frequencies, "r") as fh:
frequencies_json = json.load(fh)
if args.frequency_method == "kde":
kde_frequencies = TreeKdeFrequencies.from_json(frequencies_json)
frequencies = kde_frequencies.frequencies
# Load clades.
with open(args.clades, "r") as fh:
clades_json = json.load(fh)
clades_by_node = {
key: value["clade_membership"]
for key, value in clades_json["nodes"].items()
}
# Calculate the total frequency per clade at the most recent timepoint and
# requested timepoint in the past using non-zero tip frequencies.
current_clade_frequencies = defaultdict(float)
previous_clade_frequencies = defaultdict(float)