import os from collections import defaultdict from pathlib import Path import anndata as ad import numpy as np import pandas as pd import seaborn as sns from matplotlib import pyplot as plt from matplotlib.animation import FuncAnimation from tqdm import tqdm from hyperbolicTSNE import Datasets def plot_poincare(points, labels=None): fig, ax = plt.subplots() ax.scatter(points[:, 0], points[:, 1], c=labels, marker=".", cmap="tab10") ax.add_patch(plt.Circle((0, 0), radius=1, edgecolor="b", facecolor="None")) ax.axis("square") return fig color_dict = defaultdict(lambda: "tab10") color_dict[Datasets.C_ELEGANS] = {'ABarpaaa_lineage': '#91003f', # embryonic lineage 'Germline': '#7f2704', # Somatic gonad precursor cell 'Z1_Z4': '#800026', # Two embryonic hypodermal cells that may provide a scaffold for the early organization of ventral bodywall muscles 'XXX': '#fb8072', 'Ciliated_amphid_neuron': '#c51b8a', 'Ciliated_non_amphid_neuron': '#fa9fb5', # immune 'Coelomocyte': '#ffff33', 'T': '#54278f', # Exceratory 'Excretory_cell': '#004529', 'Excretory_cell_parent': '#006837', 'Excretory_duct_and_pore': '#238443', 'Parent_of_exc_duct_pore_DB_1_3': '#41ab5d', 'Excretory_gland': '#78c679', 'Parent_of_exc_gland_AVK': '#addd8e', 'Rectal_cell': '#d9f0a3', 'Rectal_gland': '#f7fcb9', 'Intestine': '#7fcdbb', # esophagus, crop, gizzard (usually) and intestine 'Pharyngeal_gland': '#fed976', 'Pharyngeal_intestinal_valve': '#feb24c', 'Pharyngeal_marginal_cell': '#fd8d3c', 'Pharyngeal_muscle': '#fc4e2a', 'Pharyngeal_neuron': '#e31a1c', # hypodermis (epithelial) 'Parent_of_hyp1V_and_ant_arc_V': '#a8ddb5', 'hyp1V_and_ant_arc_V': '#ccebc5', 'Hypodermis': '#253494', 'Seam_cell': '#225ea8', 'Arcade_cell': '#1d91c0', # set of six cells that form a thin cylindrical sheet between pharynx and ring neuropile 'GLR': '#1f78b4', # Glia, also called glial cells or neuroglia, are non-neuronal cells in the central nervous system 'Glia': '#377eb8', # head mesodermal cell: the middle layer of cells or tissues of an embryo 'Body_wall_muscle': '#9e9ac8', 'hmc': '#54278f', 'hmc_and_homolog': '#02818a', 'hmc_homolog': '#bcbddc', 'Intestinal_and_rectal_muscle': '#41b6c4', # Postembryonic mesoblast: the mesoderm of an embryo in its earliest stages. 'M_cell': '#3f007d', # pharyngeal gland cel 'G2_and_W_blasts': '#abdda4', 'unannotated': '#969696', 'not provided': '#969696'} color_dict[Datasets.PLANARIA] = {'neoblast 1': '#CCCCCC', 'neoblast 2': '#7f7f7f', 'neoblast 3': '#E6E6E6', 'neoblast 4': '#D6D6D6', 'neoblast 5': '#C7C7C7', 'neoblast 6': '#B8B8B8', 'neoblast 7': '#A8A8A8', 'neoblast 8': '#999999', 'neoblast 9': '#8A8A8A', 'neoblast 10': '#7A7A7A', 'neoblast 11': '#6B6B6B', 'neoblast 12': '#5C5C5C', 'neoblast 13': '#4D4D4D', 'epidermis DVb neoblast': 'lightsteelblue', 'pharynx cell type progenitors': 'slategray', 'spp-11+ neurons': '#CC4C02', 'npp-18+ neurons': '#EC7014', 'otf+ cells 1': '#993404', 'ChAT neurons 1': '#FEC44F', 'neural progenitors': '#FFF7BC', 'otf+ cells 2': '#662506', 'cav-1+ neurons': '#eec900', 'GABA neurons': '#FEE391', 'ChAT neurons 2': '#FE9929', 'muscle body': 'firebrick', 'muscle pharynx': '#CD5C5C', 'muscle progenitors': '#FF6347', 'secretory 1': 'mediumpurple', 'secretory 3': 'purple', 'secretory 4': '#CBC9E2', 'secretory 2': '#551a8b', 'early epidermal progenitors': '#9ECAE1', 'epidermal neoblasts': '#C6DBEF', 'activated early epidermal progenitors': 'lightblue', 'late epidermal progenitors 2': '#4292C6', 'late epidermal progenitors 1': '#6BAED6', 'epidermis DVb': 'dodgerblue', 'epidermis': '#2171B5', 'pharynx cell type': 'royalblue', 'protonephridia': 'pink', 'ldlrr-1+ parenchymal cells': '#d02090', 'phagocytes': 'forestgreen', 'aqp+ parenchymal cells': '#cd96cd', 'pigment': '#cd6889', 'pgrn+ parenchymal cells': 'mediumorchid', 'psap+ parenchymal cells': 'deeppink', 'glia': '#cd69c9', 'goblet cells': 'yellow', 'parenchymal progenitors': 'hotpink', 'psd+ cells': 'darkolivegreen', 'gut progenitors': 'limegreen', 'branchNe': '#4292c6', 'neutrophil': '#08306b', 'branchMo': '#9e9ac8', 'monocyte': '#54278f', 'branchEr': '#fc9272', 'erythrocyt': '#cb181d', 'megakaryoc': '#006d2c', 'branchMe': '#74c476', 'proghead': '#525252', 'interpolation': '#525252', 'Eryth': '#1F77B4', 'Gran': '#FF7F0E', 'HSPC-1': '#2CA02C', 'HSPC-2': '#D62728', 'MDP': '#9467BD', 'Meg': '#8C564B', 'Mono': '#E377C2', 'Multi-Lin': '#BCBD22', 'Myelocyte': '#17BECF', '12Baso': '#0570b0', '13Baso': '#034e7b', '11DC': '#ffff33', '18Eos': '#2CA02C', '1Ery': '#fed976', '2Ery': '#feb24c', '3Ery': '#fd8d3c', '4Ery': '#fc4e2a', '5Ery': '#e31a1c', '6Ery': '#b10026', '9GMP': '#999999', '10GMP': '#4d4d4d', '19Lymph': '#35978f', '7MEP': '#E377C2', '8Mk': '#BCBD22', '14Mo': '#4eb3d3', '15Mo': '#7bccc4', '16Neu': '#6a51a3', '17Neu': '#3f007d', 'root': '#000000'} color_dict[Datasets.MYELOID] = {'branchNe': '#4292c6', 'neutrophil': '#08306b', 'branchMo': '#9e9ac8', 'monocyte': '#54278f', 'branchEr': '#fc9272', 'erythrocyt': '#cb181d', 'megakaryoc': '#006d2c', 'branchMe': '#74c476', 'proghead': '#525252', 'root': '#000000', 'interpolation': '#bdbdbd'} legend_bbox_dict = defaultdict(lambda: (1., 1.)) legend_bbox_dict[Datasets.MNIST] = (1., 0.5) legend_bbox_dict[Datasets.PLANARIA] = (1.5, 0.5) legend_bbox_dict[Datasets.MYELOID] = (1.1, 0.5) legend_bbox_dict[Datasets.C_ELEGANS] = (1.3, 0.5) def myeloid_labels(data_home = Path("../datasets/")): full_path = Path.joinpath(data_home, "myeloid-progenitors") X = np.loadtxt(str(Path.joinpath(full_path, "MyeloidProgenitors.csv")), delimiter=",", skiprows=1, usecols=np.arange(11)) return np.loadtxt(str(Path.joinpath(full_path, "MyeloidProgenitors.csv")), delimiter=",", skiprows=1, usecols=11, dtype=str) def c_elegans_labels(data_home = Path("../datasets/")): full_path = Path.joinpath(data_home, "c_elegans") ad_obj = ad.read_h5ad(str(Path.joinpath(full_path, "packer2019.h5ad"))) X = ad_obj.X return np.array(ad_obj.obs.cell_type) def planaria_labels(data_home = Path("../datasets/")): full_path = Path.joinpath(data_home, "planaria") return np.loadtxt(str(Path.joinpath(full_path, "R_annotation.txt")), delimiter=",", dtype=str) def mnist_labels(data_home = Path("../datasets/")): full_path = Path.joinpath(data_home, 'mnist') labels_path_train = Path.joinpath(full_path, 'train-labels-idx1-ubyte.gz') labels_path_test = Path.joinpath(full_path, 't10k-labels-idx1-ubyte.gz') labels_arr = [] import gzip with gzip.open(labels_path_train, 'rb') as lbpath: br = np.frombuffer(lbpath.read(), dtype=np.uint8, offset=8) labels_arr.append(br) with gzip.open(labels_path_test, 'rb') as lbpath: br = np.frombuffer(lbpath.read(), dtype=np.uint8, offset=8) labels_arr.append(br) labels = np.concatenate(labels_arr, axis=0) return list(map(str, labels)) # labels_dict = defaultdict(lambda: None) # labels_dict[Datasets.MNIST] = mnist_labels() # labels_dict[Datasets.PLANARIA] = planaria_labels() # labels_dict[Datasets.C_ELEGANS] = c_elegans_labels() # labels_dict[Datasets.MYELOID] = myeloid_labels() def save_poincare_teaser(points, file_name, str_labels=None, dataset=None, save_fig_kwargs=dict()): df = pd.DataFrame({"x": points[:, 0], "y": points[:, 1]}) fig, ax = plt.subplots() point_size = 2 font_size = 5 alpha = 1.0 if str_labels is None: str_labels = labels_dict[dataset] sns.scatterplot(data=df, x="x", y="y", hue=str_labels, hue_order=np.unique(str_labels), palette=color_dict[dataset], alpha=alpha, edgecolor="none", ax=ax, s=point_size) lgd = ax.legend(fontsize=font_size, loc='right', bbox_to_anchor=legend_bbox_dict[dataset], facecolor='white', frameon=False, ncol=2 if dataset is Datasets.PLANARIA else 1) circle = plt.Circle((0, 0), radius=1, fc='none', color='black') ax.add_patch(circle) ax.plot(0, 0, '.', c=(0, 0, 0), ms=4) fig.tight_layout() ax.axis('off') ax.axis('equal') ax.set_ylim([-1.01, 1.01]) ax.set_xlim([-1.01, 1.01]) plt.tight_layout() plt.savefig(file_name, bbox_extra_artists=(lgd,), bbox_inches='tight', **save_fig_kwargs) def plot_poincare_zoomed(points, labels=None): fig, ax = plt.subplots() ax.scatter(points[:, 0], points[:, 1], c=labels, marker=".", cmap="tab10") ax.axis("square") ax.add_patch(plt.Circle((0, 0), radius=1, edgecolor="b", facecolor="None")) return fig def animate(log_dict, labels, file_name, fast=False, is_hyperbolic=True, plot_ee=False, first_frame=None): scatter_data = [] if first_frame is None else [("-1", first_frame)] for subdir, dirs, files in os.walk(log_dict["log_path"]): for fi, file in enumerate(sorted(files, key=lambda x: int(x.split(", ")[0]))): root, ext = os.path.splitext(file) if ext == ".csv": total_file = subdir.replace("\\", "/") + "/" + file if (not fast or fi % 10 == 0) and (plot_ee or subdir.split("/")[-1].endswith("1")): data = np.genfromtxt(total_file, delimiter=',') scatter_data.append((str(fi), data)) fig, ax = plt.subplots() _, data = scatter_data[0 if is_hyperbolic else -1] scatter = ax.scatter(data[:, 0], data[:, 1], c=labels, marker=".", linewidth=0.5, s=20, cmap="tab10") if is_hyperbolic: uc = plt.Circle((0, 0), radius=1, edgecolor="b", facecolor="None") ax.add_patch(uc) ax.axis("square") print("Animation being saved to: " + file_name) pbar = tqdm(desc="Animating: ", total=len(scatter_data)) def update(i): pbar.update() sd = scatter_data[i] scatter.set_offsets(sd[1]) ax.set_title(f'Scatter (epoch {sd[0]})') return scatter, anim = FuncAnimation(fig, update, frames=len(scatter_data), interval=50, blit=True, save_count=50) anim.save(file_name) plt.clf() plt.close() pbar.close()