Allow axis passing to sample_subgraph

[hyanwong]

And return the graph and pos for further annotation if needed.

For example, we can then adjust the aspect ratio:

fig, ax = plt.subplots(1, 1, figsize=(40, 20))

sc2ts.sample_subgraph(
    712029, 
    ts, 
    treeinfo, 
    mutations_json_filepath="../data/consensus_mutations.json",
    imputation_source='Imputed_lineage',
    ax=ax,
)