run.py

import numpy as np
import cmapy
import cv2
import scipy.ndimage as ndim
import tqdm
import argparse
import imageio as io
import os
import json
import skimage.measure
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 18})
import sys
sys.path.append("./build")
import pycellpotts


def correct_spins(tif):
    cells = np.unique(tif)
    tif_copy = tif.copy()
    cell_list = list(np.sort(np.unique(tif)))
    print(cell_list)
    tmp = np.zeros_like(tif)
    print(tmp.shape)
    tmp[0, :], tmp[:, 0], tmp[-1], tmp[:, -1] = 1, 1, 1, 1

    for c in cells:
        cell = tif == c
        labeled, num = skimage.measure.label(cell, return_num=True)

        if num > 1:
            renumber = []
            border = []
            border_pxs = {}
            for l in np.unique(labeled):
                if l == 0:
                    continue
                borders_masked = tmp * (labeled == l)

                if np.sum(borders_masked) > 0:
                    border_pxs[l] = borders_masked
                    border.append(l)
                else:
                    # automatically new cell
                    renumber.append(l)
            if len(border) == 1:
                renumber.append(border[0])

            matched_pairs = {}
            for b in border:
                matches = False
                if b in [item for sublist in matched_pairs.values() for item in sublist]:
                    continue
                if b not in matched_pairs.keys():
                    matched_pairs[b] = []
                for b2 in border:
                    if b == b2:
                        continue

                    b_pxs, b2_pxs = border_pxs[b], border_pxs[b2]

                    if np.sum(b_pxs[0]) > 0 or np.sum(b_pxs[-1]) > 0:
                        flipped_up = np.flipud(b_pxs)
                        dilated = ndim.binary_dilation(flipped_up)
                        overlap = dilated * b2_pxs
                        if np.sum(overlap) > 0:
                            matches = True
                            matched_pairs[b].append(b2)
                    if np.sum(b_pxs[:, 0]) > 0 or np.sum(b_pxs[:, -1]) > 0:
                        flipped_lr = np.fliplr(b_pxs)
                        dilated = ndim.binary_dilation(flipped_lr)
                        overlap = dilated * b2_pxs
                        if np.sum(overlap) > 0:
                            matches = True
                            matched_pairs[b].append(b2)
                    # TODO  figure out if it happens any other time
                    if (np.sum(b_pxs[0]) > 0 or np.sum(b_pxs[-1]) > 0) and (
                            np.sum(b_pxs[:, 0]) > 0 or np.sum(b_pxs[:, -1]) > 0):
                        flipped_up = np.flipud(b_pxs)
                        flipped_both = np.fliplr(flipped_up)
                        dilated = ndim.binary_dilation(flipped_both)
                        overlap = dilated * b2_pxs
                        if np.sum(overlap) > 0:
                            matches = True
                            matched_pairs[b].append(b2)

                if not matches:
                    if b not in renumber:
                        renumber.append(b)

            if len(matched_pairs.keys()) > 1 or len(renumber) > 0:
                for b in matched_pairs.keys():
                    for m in matched_pairs[b]:
                        labeled[labeled == m] = b
                    # if c == 308:
                    #     _, axs = plt.subplots(1, 2)
                    #     axs[0].imshow(labeled)
                    #     axs[1].imshow(labeled == b)
                    #     plt.show()
                    if b not in renumber:
                        renumber.append(b)
            # print(matched_pairs.keys())

            if len(renumber) > 0:
                print(renumber)
                for r in renumber:
                    _, axs = plt.subplots(1, 2)
                    axs[0].imshow(labeled)
                    axs[1].imshow(labeled == r)
                    plt.show()
                for r in renumber[1:]:
                    new_spin = cell_list[-1] + 1
                    tif_copy[labeled == r] = new_spin
                    cell_list.append(new_spin)
                    plt.imshow(tif_copy)
                    plt.show()
    return tif_copy


def run_cpm(nr_updates=500, beta=1., length=6, steps=2000, q_max=400, q_initial=20, height=600, width=600, id_area_init=60, lamb=1,
            mu=.02, wound_area=15, divide=True, filled=False, j_same_cell=0, j_diff_cell=-5, j_medium_cell=-10, save=False, show=False):
    if save:
        all_args = locals()
        path = 'results'
        os.makedirs(f'{path}/gifs', exist_ok=True)
        os.makedirs(f'{path}/plots', exist_ok=True)
        os.makedirs(f'{path}/args', exist_ok=True)

        i = 0
        while os.path.exists(f"{path}/gifs/exp_{i}.gif"):
            i += 1
        file = f'exp_{i}'

    areas = []
    video = []
    wound_time = 0

    if show:
        cv2.namedWindow('lattice', cv2.WINDOW_NORMAL)
    if mu > 0.02 or j_medium_cell < -10:
        q_max = 5000

    cpm = pycellpotts.CPM(b=beta, l=length, p=steps, q=q_max, q1=q_initial, h=height, w=width, tip_area=id_area_init,
                          lam=lamb, prolif_prob=mu, wa=wound_area, div=divide, fill=filled, j_same=j_same_cell,
                          j_diff=j_diff_cell, j_med_cell=j_medium_cell)
    for i in tqdm.tqdm(range(nr_updates)):
        cpm.full_update()
        lattice = cpm.lattice_as_numpy()

        if cpm.wounded and wound_time == 0:
            wound_time = i - 1

        med_a = np.sum(lattice == 0)
        cpm_areas = cpm.areas_as_numpy()
        cell_a = np.mean(cpm_areas[cpm_areas != 0])
        areas.append([med_a, cell_a])

        if i == 90:
            io.imsave(f'/home/lola/Workspace/try.tif', lattice)
        lattice_cm = cv2.applyColorMap(lattice.astype(np.uint8), cmapy.cmap('gist_ncar_r'))
        if show:
            cv2.imshow('lattice', lattice_cm)
            cv2.waitKey(1)
        video.append(lattice_cm)

        if med_a == cpm.height * cpm.width:
            break
        if med_a == 0 and cpm.wounded:
            break
        # if i == 20:
        #     for g in range(len(cpm.lattice)):
        #         cpm.lattice[g] = 0

    areas = np.array(areas)
    print(areas.shape)
    if len(areas) > 0 and (save or show):
        areas = np.array(areas)
        fig, axs = plt.subplots(1, 2, figsize=(12, 8))
        axs[0].plot(areas[wound_time:, 0], 'b')
        axs[1].plot(areas[:, 1], 'g')

        axs[0].set_ylim(0, areas[wound_time:, 0].max() + 30)
        axs[1].set_ylim(0, areas[:, 1].max() + 30)
        axs[0].set_xlim(0, len(areas[wound_time:]) + 10)
        axs[1].set_xlim(0, len(areas) + 10)
        axs[0].set_xlabel("Time")
        axs[0].set_ylabel("Wound area")
        axs[1].set_xlabel("Time")
        axs[1].set_ylabel("Average cell area")
        axs[1].axvline(wound_time+1, ls='--', color='gray')
        fig.tight_layout()
        # plt.axvline(wound_time + 1, '--', color='gray')
        if save:
            plt.savefig(f'{path}/plots/{file}.png', )
        else:
            plt.show()
        plt.close()
    if save:
        video = np.array(video)
        io.mimwrite(f'{path}/gifs/{file}.gif', video, fps=30)
        with open(f'{path}/args/{file}.json', 'w') as file:
            file.write(json.dumps(all_args))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--nr_updates", dest="nr_updates", type=int, help="Number of all frame updates",
                        default=1000)
    parser.add_argument("-b", "--beta", dest="beta", type=float, help="Temperature of Monte-Carlo step", default=1)
    parser.add_argument("-l", "--length", dest="length", type=int, help="length of the sketch", default=1)
    parser.add_argument("-p", "--steps", dest="steps", type=int, help="Steps of spin update per frame update",
                        default=120*120)
    parser.add_argument("-q", "--q_max", dest="q_max", type=int, help="Max number of cells", default=2500)
    parser.add_argument("-q1", "--q_initial", dest="q_initial", type=int, help="Initial number of cells", default=100)
    parser.add_argument("-he", "--height", dest="height", type=int, help="Height of sketch", default=120)
    parser.add_argument("-w", "--width", dest="width", type=int, help="Width of sketch", default=120)
    parser.add_argument("-i", "--id_area_init", dest="id_area_init", type=int, help="Ideal area", default=20)
    parser.add_argument("-la", "--lamb", dest="lamb", type=float, help="Area constraint strength", default=1.)
    parser.add_argument("-mu", "--proliferation_probability", dest="mu", type=float, help="Proliferation probability",
                        default=.02)
    parser.add_argument("-wa", "--wound_area", dest="wound_area", type=int, help="Wounded area", default=15)
    parser.add_argument("-div", "--divide", dest="divide", type=bool, help="Proliferation phase", default=True)
    parser.add_argument("-f", "--filled", dest="filled", type=bool, help="Filled with cells", default=False)
    parser.add_argument("-c", "--j_same_cell", dest="j_same_cell", type=int, help="Adhesion of same cell", default=0)
    parser.add_argument("-d", "--j_diff_cell", dest="j_diff_cell", type=int, help="Adhesion of different cells",
                        default=-2)
    parser.add_argument("-m", "--j_medium_cell", dest="j_medium_cell", type=int, help="Adhesion medium and cell",
                        default=-5)
    parser.add_argument("-s", "--save", dest="save", type=bool, help="Save results", default=False)
    parser.add_argument("-v", "--show", dest="show", type=bool, help="Show results", default=True)

    arguments = parser.parse_args()
    run_cpm(**vars(arguments))