def data_to_tensor_tree(tree_test_data: TreeTestData) -> TensorflowRootedTree: numpy_tree = NumpyRootedTree( node_heights=tree_test_data.node_heights, sampling_times=tree_test_data.sampling_times, parent_indices=tree_test_data.parent_indices, ) tf_tree = convert_tree_to_tensor(numpy_tree) return tf_tree
def test_BirthDeathContemporarySampling_log_prob(): newick = "((((human:0.024003,(chimp:0.010772,bonobo:0.010772):0.013231):0.012035,gorilla:0.036038):0.033087000000000005,orangutan:0.069125):0.030456999999999998,siamang:0.099582);" tree = convert_tree_to_tensor(parse_newick(newick)) birth_diff_rate = tf.constant(1.0, dtype=DEFAULT_FLOAT_DTYPE_TF) relative_death_rate = tf.constant(0.5, dtype=DEFAULT_FLOAT_DTYPE_TF) expected = 1.2661341104158121 # From BEAST 2 BirthDeathGerhard08ModelTest dist = BirthDeathContemporarySampling(tree.taxon_count, birth_diff_rate, relative_death_rate) res = dist.log_prob(tree) assert_allclose(res.numpy(), expected)
def test_sample_weighted(wnv_newick_file, wnv_fasta_file, hky_params): alignment = Alignment(wnv_fasta_file) numpy_tree = parse_newick(wnv_newick_file) tensor_tree = convert_tree_to_tensor(numpy_tree) transition_prob_tree = get_transition_probabilities_tree( tensor_tree.get_unrooted_tree(), HKY(), **hky_params) log_prob_uncompressed = compute_log_prob_uncompressed( alignment, transition_prob_tree, hky_params["frequencies"]) log_prob_compressed = compute_log_prob_compressed( alignment, transition_prob_tree, hky_params["frequencies"]) assert_allclose(log_prob_uncompressed.numpy(), log_prob_compressed.numpy())
def transition_prob_tree(flat_tree_test_data): tree = convert_tree_to_tensor( NumpyRootedTree( heights=flat_tree_test_data.heights, parent_indices=flat_tree_test_data.parent_indices, ) ).get_unrooted_tree() state_count = 5 transition_probs = tf.fill( tree.branch_lengths.shape + (state_count, state_count), 1.0 / state_count ) return TensorflowUnrootedTree( branch_lengths=transition_probs, topology=numpy_topology_to_tensor(tree.topology), )
def test_log_prob_conditioned_hky(hky_params, newick_fasta_file_dated): from treeflow.acceleration.bito.beagle import ( phylogenetic_likelihood as beagle_likelihood, ) newick_file, fasta_file, dated = newick_fasta_file_dated subst_model = HKY() tensor_tree = convert_tree_to_tensor( parse_newick(newick_file)).get_unrooted_tree() alignment = Alignment(fasta_file) sequences = alignment.get_encoded_sequence_tensor(tensor_tree.taxon_set) treeflow_func = lambda blens: Sample( LeafCTMC( get_transition_probabilities_tree( tensor_tree.with_branch_lengths(blens), subst_model, ** hky_params), hky_params["frequencies"], ), sample_shape=alignment.site_count, ).log_prob(sequences) beagle_func, _ = beagle_likelihood(fasta_file, subst_model, newick_file=newick_file, dated=dated, **hky_params) blens = tensor_tree.branch_lengths with tf.GradientTape() as tf_t: tf_t.watch(blens) tf_ll = treeflow_func(blens) tf_gradient = tf_t.gradient(tf_ll, blens) with tf.GradientTape() as bito_t: bito_t.watch(blens) bito_ll = beagle_func(blens) bito_gradient = bito_t.gradient(bito_ll, blens) assert_allclose(tf_ll.numpy(), bito_ll.numpy()) assert_allclose(tf_gradient.numpy(), bito_gradient.numpy())
def tree_from_ratio_test_data( ratio_test_data: RatioTestData) -> TensorflowRootedTree: return convert_tree_to_tensor( numpy_tree_from_ratio_test_data(ratio_test_data), )
def test_get_anchor_heights_tensor(ratio_test_data: RatioTestData): tree = convert_tree_to_tensor( numpy_tree_from_ratio_test_data(ratio_test_data)) res = get_anchor_heights_tensor(tree.topology, tree.sampling_times) assert_allclose(res.numpy(), ratio_test_data.anchor_heights)
def treeflow_vi( input, topology, num_steps, optimizer, substitution_model, clock_model, tree_prior, learning_rate, init_values, approx_output, trace_output, ): optimizer = optimizer_classes[optimizer](learning_rate=learning_rate) print(f"Parsing topology {topology}") tree = convert_tree_to_tensor(parse_newick(topology)) print(f"Parsing alignment {input}") alignment = Alignment(input).get_compressed_alignment() encoded_sequences = alignment.get_encoded_sequence_tensor(tree.taxon_set) pattern_counts = alignment.get_weights_tensor() print(f"Parsing initial values...") init_values_dict = (None if init_values is None else { key: tf.constant(value, dtype=DEFAULT_FLOAT_DTYPE_TF) for key, value in parse_init_values(init_values).items() }) if (clock_model == _FIXED_STRICT and substitution_model == _JC_KEY and tree_prior == _CONSTANT_COALESCENT): model = get_example_phylo_model( taxon_count=tree.taxon_count, site_count=alignment.site_count, sampling_times=tree.sampling_times, pattern_counts=pattern_counts, init_values=init_values_dict, ) else: raise ValueError( "Only example with JC, fixed-rate strict clock, constant coalescent supported for now" ) pinned_model = model.experimental_pin(alignment=encoded_sequences) model_names = set(pinned_model._flat_resolve_names()) init_loc = { key: value for key, value in init_values_dict.items() if key in model_names } init_loc["tree"] = tree print(f"Running VI for {num_steps} iterations...") approx, trace = fit_fixed_topology_variational_approximation( model=pinned_model, topologies={DEFAULT_TREE_NAME: tree.topology}, init_loc=init_loc, optimizer=optimizer, num_steps=num_steps, ) print("Inference complete") print("Approx sample:") print(tf.nest.map_structure(lambda x: x.numpy(), approx.sample())) if approx_output is not None: print(f"Saving approximation to {approx_output}...") with open(approx_output, "wb") as f: pickle.dump(approx, f) # TODO: Support saving approximation if trace_output is not None: print(f"Saving trace to {trace_output}...") with open(trace_output, "wb") as f: pickle.dump(trace, f) print("Exiting...")
def treeflow_benchmark(input: str, tree: str, replicates: int, output: tp.Optional[io.StringIO], dtype: str, scaler: float, precompile: bool, eager: bool, memory: bool, use_bito: bool): print("Parsing input...") alignment = Alignment(input) numpy_tree = parse_newick(tree, remove_zero_edges=True) tf_dtype = DTYPE_MAPPING[dtype] tensor_tree = convert_tree_to_tensor(numpy_tree, height_dtype=tf_dtype) scaler_tensor = tf.constant(scaler, dtype=tf_dtype) computations = [ "treelikelihood", "ratio_transform", "ratio_transform_jacobian", "constant_coalescent", ] tasks = ["gradient", "evaluation"] jits = [True, False] if eager else [True] if output: output.write("function,mode,JIT,time,logprob") if memory: output.write(",max_mem") output.write("\n") print("Starting benchmark...") if use_bito: from treeflow.acceleration.bito.instance import get_instance dated = not np.allclose(numpy_tree.sampling_times, 0.0) bito_instance = get_instance(tree, dated=dated) else: bito_instance = None for (computation, task, jit) in product(computations, tasks, jits): print( f"Benchmarking {computation} {task}{' in function mode' if jit else ''}..." ) benchmark_args = ( tensor_tree, alignment, tf_dtype, scaler_tensor, computation, task, jit, precompile, replicates, ) max_mem: tp.Optional[float] if memory: from memory_profiler import memory_usage max_mem, (time, value) = memory_usage( (benchmark, benchmark_args, dict(bito_instance=bito_instance)), retval=True, max_usage=True, max_iterations=1, ) else: time, value = benchmark(*benchmark_args, bito_instance=bito_instance) max_mem = None if output: jit_str = "on" if jit else "off" output.write( f"{computation},{task},{jit_str},{time},{value if value else ''}" ) if max_mem is not None: output.write(f",{max_mem}") output.write("\n") if max_mem is not None: print(f"Max memory usage: {max_mem}") print("\n") print("Benchmark complete")
def get_treeflow_forward_func(numpy_tree, tensor_constant): anchor_heights = tensor_constant(get_anchor_heights(numpy_tree)) return NodeHeightRatioBijector( convert_tree_to_tensor(numpy_tree).topology, anchor_heights).forward