def ratio_transform(args): replicates = args.replicates tree = read_tree(args.tree, True, True) taxa_count = len(tree.taxon_namespace) taxa = [] for node in tree.leaf_node_iter(): taxa.append(Taxon(node.label, {'date': node.date})) ratios_root_height = Parameter( "internal_heights", torch.tensor([0.5] * (taxa_count - 2) + [10]) ) tree_model = ReparameterizedTimeTreeModel( "tree", tree, Taxa('taxa', taxa), ratios_root_height ) ratios_root_height.tensor = tree_model.transform.inv( heights_from_branch_lengths(tree) ) @benchmark def fn(ratios_root_height): return tree_model.transform( ratios_root_height, ) @benchmark def fn_grad(ratios_root_height): heights = tree_model.transform( ratios_root_height, ) heights.backward(torch.ones_like(ratios_root_height)) ratios_root_height.grad.data.zero_() return heights total_time, heights = fn(args.replicates, ratios_root_height.tensor) print(f' {replicates} evaluations: {total_time}') ratios_root_height.requires_grad = True grad_total_time, heights = fn_grad(args.replicates, ratios_root_height.tensor) print(f' {replicates} gradient evaluations: {grad_total_time}') if args.output: args.output.write(f"ratio_transform,evaluation,off,{total_time},\n") args.output.write(f"ratio_transform,gradient,off,{grad_total_time},\n") print(' JIT off') @benchmark def fn2(ratios_root_height): return transform( tree_model.transform._forward_indices, tree_model.transform._bounds, ratios_root_height, ) @benchmark def fn2_grad(ratios_root_height): heights = transform( tree_model.transform._forward_indices, tree_model.transform._bounds, ratios_root_height, ) heights.backward(torch.ones_like(ratios_root_height)) ratios_root_height.grad.data.zero_() return heights ratios_root_height.requires_grad = False total_time, heights = fn2(args.replicates, ratios_root_height.tensor) print(f' {replicates} evaluations: {total_time}') ratios_root_height.requires_grad = True total_time, heights = fn2_grad(args.replicates, ratios_root_height.tensor) print(f' {replicates} gradient evaluations: {total_time}') print(' JIT on') transform_script = torch.jit.script(transform) @benchmark def fn2_jit(ratios_root_height): return transform_script( tree_model.transform._forward_indices, tree_model.transform._bounds, ratios_root_height, ) @benchmark def fn2_grad_jit(ratios_root_height): heights = transform_script( tree_model.transform._forward_indices, tree_model.transform._bounds, ratios_root_height, ) heights.backward(torch.ones_like(ratios_root_height)) ratios_root_height.grad.data.zero_() return heights ratios_root_height.requires_grad = False total_time, heights = fn2_jit(args.replicates, ratios_root_height.tensor) print(f' {replicates} evaluations: {total_time}') ratios_root_height.requires_grad = True total_time, heights = fn2_grad_jit(args.replicates, ratios_root_height.tensor) print(f' {replicates} gradient evaluations: {total_time}') print('ratio_transform v2 JIT off') @benchmark def fn3(ratios_root_height): return transform2( tree_model.transform._forward_indices.tolist(), tree_model.transform._bounds, ratios_root_height, ) @benchmark def fn3_grad(ratios_root_height): heights = transform2( tree_model.transform._forward_indices.tolist(), tree_model.transform._bounds, ratios_root_height, ) heights.backward(torch.ones_like(ratios_root_height)) ratios_root_height.grad.data.zero_() return heights ratios_root_height.requires_grad = False total_time, heights = fn3(args.replicates, ratios_root_height.tensor) print(f' {replicates} evaluations: {total_time}') ratios_root_height.requires_grad = True total_time, heights = fn3_grad(args.replicates, ratios_root_height.tensor) print(f' {replicates} gradient evaluations: {total_time}') print('ratio_transform v2 JIT on') transform2_script = torch.jit.script(transform2) @benchmark def fn3_jit(ratios_root_height): return transform2_script( tree_model.transform._forward_indices.tolist(), tree_model.transform._bounds, ratios_root_height, ) @benchmark def fn3_grad_jit(ratios_root_height): heights = transform2_script( tree_model.transform._forward_indices.tolist(), tree_model.transform._bounds, ratios_root_height, ) heights.backward(torch.ones_like(ratios_root_height)) ratios_root_height.grad.data.zero_() return heights ratios_root_height.requires_grad = False total_time, heights = fn3_jit(args.replicates, ratios_root_height.tensor) print(f' {replicates} evaluations: {total_time}') ratios_root_height.requires_grad = True total_time, heights = fn3_grad_jit(args.replicates, ratios_root_height.tensor) print(f' {replicates} gradient evaluations: {total_time}')
def constant_coalescent(args): tree = read_tree(args.tree, True, True) taxa_count = len(tree.taxon_namespace) taxa = [] for node in tree.leaf_node_iter(): taxa.append(Taxon(node.label, {'date': node.date})) ratios_root_height = Parameter( "internal_heights", torch.tensor([0.5] * (taxa_count - 2) + [20.0]) ) tree_model = TimeTreeModel("tree", tree, Taxa('taxa', taxa), ratios_root_height) tree_model._internal_heights.tensor = heights_from_branch_lengths(tree) pop_size = torch.tensor([4.0]) print('JIT off') @benchmark def fn(tree_model, pop_size): return ConstantCoalescent(pop_size).log_prob(tree_model.node_heights) @benchmark def fn_grad(tree_model, pop_size): log_p = ConstantCoalescent(pop_size).log_prob(tree_model.node_heights) log_p.backward() ratios_root_height.tensor.grad.data.zero_() pop_size.grad.data.zero_() return log_p total_time, log_p = fn(args.replicates, tree_model, pop_size) print(f' {args.replicates} evaluations: {total_time} {log_p}') ratios_root_height.requires_grad = True pop_size.requires_grad_(True) grad_total_time, grad_log_p = fn_grad(args.replicates, tree_model, pop_size) print(f' {args.replicates} gradient evaluations: {grad_total_time}') if args.output: args.output.write( f"coalescent,evaluation,off,{total_time},{log_p.squeeze().item()}\n" ) args.output.write( f"coalescent,gradient,off,{grad_total_time},{grad_log_p.squeeze().item()}\n" ) if args.debug: tree_model.heights_need_update = True log_p = ConstantCoalescent(pop_size).log_prob(tree_model.node_heights) log_p.backward() print('gradient ratios: ', ratios_root_height.grad) print('gradient pop size: ', pop_size.grad) ratios_root_height.tensor.grad.data.zero_() pop_size.grad.data.zero_() print('JIT on') log_prob_script = torch.jit.script(log_prob) @benchmark def fn_jit(tree_model, pop_size): return log_prob_script(tree_model.node_heights, pop_size) @benchmark def fn_grad_jit(tree_model, pop_size): log_p = log_prob_script(tree_model.node_heights, pop_size) log_p.backward() ratios_root_height.tensor.grad.data.zero_() pop_size.grad.data.zero_() return log_p ratios_root_height.requires_grad = False pop_size.requires_grad_(False) total_time, log_p = fn_jit(args.replicates, tree_model, pop_size) print(f' {args.replicates} evaluations: {total_time} {log_p}') ratios_root_height.requires_grad = True pop_size.requires_grad_(True) grad_total_time, grad_log_p = fn_grad_jit(args.replicates, tree_model, pop_size) print(f' {args.replicates} gradient evaluations: {grad_total_time}') if args.output: args.output.write( f"coalescent,evaluation,on,{total_time},{log_p.squeeze().item()}\n" ) args.output.write( f"coalescent,gradient,on,{grad_total_time},{grad_log_p.squeeze().item()}\n" ) if args.all: print('make sampling times unique and count them:') @benchmark def fn3(tree_model, pop_size): tree_model.heights_need_update = True node_heights = torch.cat( (x, tree_model.node_heights[..., tree_model.taxa_count :]) ) return log_prob_squashed( pop_size, node_heights, counts, tree_model.taxa_count ) @benchmark def fn3_grad(tree_model, ratios_root_height, pop_size): tree_model.heights_need_update = True node_heights = torch.cat( (x, tree_model.node_heights[..., tree_model.taxa_count :]) ) log_p = log_prob_squashed( pop_size, node_heights, counts, tree_model.taxa_count ) log_p.backward() ratios_root_height.tensor.grad.data.zero_() pop_size.grad.data.zero_() return log_p x, counts = torch.unique(tree_model.sampling_times, return_counts=True) counts = torch.cat((counts, torch.tensor([-1] * (taxa_count - 1)))) with torch.no_grad(): total_time, log_p = fn3(args.replicates, tree_model, pop_size) print(f' {args.replicates} evaluations: {total_time} ({log_p})') total_time, log_p = fn3_grad( args.replicates, tree_model, ratios_root_height, pop_size ) print(f' {args.replicates} gradient evaluations: {total_time}')
def ratio_transform_jacobian(args): tree = read_tree(args.tree, True, True) taxa = [] for node in tree.leaf_node_iter(): taxa.append(Taxon(node.label, {'date': node.date})) taxa_count = len(taxa) ratios_root_height = Parameter( "internal_heights", torch.tensor([0.5] * (taxa_count - 1) + [20]) ) tree_model = ReparameterizedTimeTreeModel( "tree", tree, Taxa('taxa', taxa), ratios_root_height ) ratios_root_height.tensor = tree_model.transform.inv( heights_from_branch_lengths(tree) ) @benchmark def fn(ratios_root_height): internal_heights = tree_model.transform(ratios_root_height) return tree_model.transform.log_abs_det_jacobian( ratios_root_height, internal_heights ) @benchmark def fn_grad(ratios_root_height): internal_heights = tree_model.transform(ratios_root_height) log_det_jac = tree_model.transform.log_abs_det_jacobian( ratios_root_height, internal_heights ) log_det_jac.backward() ratios_root_height.grad.data.zero_() return log_det_jac print(' JIT off') total_time, log_det_jac = fn(args.replicates, ratios_root_height.tensor) print(f' {args.replicates} evaluations: {total_time} ({log_det_jac})') ratios_root_height.requires_grad = True grad_total_time, grad_log_det_jac = fn_grad( args.replicates, ratios_root_height.tensor ) print( f' {args.replicates} gradient evaluations: {grad_total_time}' f' ({grad_log_det_jac})' ) if args.output: args.output.write( f"ratio_transform_jacobian,evaluation,off,{total_time}," f"{log_det_jac.squeeze().item()}\n" ) args.output.write( f"ratio_transform_jacobian,gradient,off,{grad_total_time}," f"{grad_log_det_jac.squeeze().item()}\n" ) if args.debug: internal_heights = tree_model.transform(ratios_root_height.tensor) log_det_jac = tree_model.transform.log_abs_det_jacobian( ratios_root_height.tensor, internal_heights ) log_det_jac.backward() print(ratios_root_height.grad)