Initial commit

3 files changed, 310 insertions, 0 deletions

diff --git a/code/sunlab/common/plotting/__init__.py b/code/sunlab/common/plotting/__init__.py
new file mode 100644
index 0000000..d6873aa
--- /dev/null
+++ b/code/sunlab/common/plotting/__init__.py
@@ -0,0 +1,2 @@
+from .colors import *
+from .base import *
diff --git a/code/sunlab/common/plotting/base.py b/code/sunlab/common/plotting/base.py
new file mode 100644
index 0000000..aaf4a94
--- /dev/null
+++ b/code/sunlab/common/plotting/base.py
@@ -0,0 +1,270 @@
+from matplotlib import pyplot as plt
+
+
+def blank_plot(_plt=None, _xticks=False, _yticks=False):
+    """# Remove Plot Labels"""
+    if _plt is None:
+        _plt = plt
+    _plt.xlabel("")
+    _plt.ylabel("")
+    _plt.title("")
+    tick_params = {
+        "which": "both",
+        "bottom": _xticks,
+        "left": _yticks,
+        "right": False,
+        "labelleft": False,
+        "labelbottom": False,
+    }
+    _plt.tick_params(**tick_params)
+    for child in plt.gcf().get_children():
+        if child._label == "<colorbar>":
+            child.set_yticks([])
+    axs = _plt.gcf().get_axes()
+    try:
+        axs = axs.flatten()
+    except:
+        ...
+    for ax in axs:
+        ax.set_xlabel("")
+        ax.set_ylabel("")
+        ax.set_title("")
+        ax.tick_params(**tick_params)
+
+
+def save_plot(name, _plt=None, _xticks=False, _yticks=False, tighten=True):
+    """# Save Plot in Multiple Formats"""
+    assert type(name) == str, "Name must be string"
+    from os.path import dirname
+    from os import makedirs
+
+    makedirs(dirname(name), exist_ok=True)
+    if _plt is None:
+        from matplotlib import pyplot as plt
+        _plt = plt
+    _plt.savefig(name + ".png", dpi=1000)
+    blank_plot(_plt, _xticks=_xticks, _yticks=_yticks)
+    if tighten:
+        from matplotlib import pyplot as plt
+        plt.tight_layout()
+    _plt.savefig(name + ".pdf")
+    _plt.savefig(name + ".svg")
+
+
+def scatter_2d(data_2d, labels=None, _plt=None, **matplotlib_kwargs):
+    """# Scatter 2d Data
+
+    - data_2d: 2d-dataset to plot
+    - labels: labels for each case
+    - _plt: Optional specific plot to transform"""
+    from .colors import Pcolor
+
+    if _plt is None:
+        _plt = plt
+
+    def _filter(data, labels=None, _filter_on=None):
+        if labels is None:
+            return data, False
+        else:
+            return data[labels == _filter_on], True
+
+    for _class in [2, 3, 1, 0]:
+        local_data, has_color = _filter(data_2d, labels, _class)
+        if has_color:
+            _plt.scatter(
+                local_data[:, 0],
+                local_data[:, 1],
+                color=Pcolor[_class],
+                **matplotlib_kwargs
+            )
+        else:
+            _plt.scatter(local_data[:, 0], local_data[:, 1], **matplotlib_kwargs)
+            break
+    return _plt
+
+
+def plot_contour(two_d_mask, color="black", color_map=None, raise_error=False):
+    """# Plot Contour of Mask"""
+    from matplotlib.pyplot import contour
+    from numpy import mgrid
+
+    z = two_d_mask
+    x, y = mgrid[: z.shape[1], : z.shape[0]]
+    if color_map is not None:
+        try:
+            color = color_map(color)
+        except Exception as e:
+            if raise_error:
+                raise e
+    try:
+        contour(x, y, z.T, colors=color, linewidth=0.5)
+    except Exception as e:
+        if raise_error:
+            raise e
+
+
+def gaussian_smooth_plot(
+    xy,
+    sigma=0.1,
+    extent=[-1, 1, -1, 1],
+    grid_n=100,
+    grid=None,
+    do_normalize=True,
+):
+    """# Plot Data with Gaussian Smoothening around point"""
+    from numpy import array, ndarray, linspace, meshgrid, zeros_like
+    from numpy import pi, sqrt, exp
+    from numpy.linalg import norm
+
+    if not type(xy) == ndarray:
+        xy = array(xy)
+    if grid is not None:
+        XY = grid
+    else:
+        X = linspace(extent[0], extent[1], grid_n)
+        Y = linspace(extent[2], extent[3], grid_n)
+        XY = array(meshgrid(X, Y)).T
+    smoothed = zeros_like(XY[:, :, 0])
+    factor = 1
+    if do_normalize:
+        factor = (sqrt(2 * pi) * sigma) ** 2.
+    if len(xy.shape) == 1:
+        smoothed = exp(-((norm(xy - XY, axis=-1) / (sqrt(2) * sigma)) ** 2.0)) / factor
+    else:
+        try:
+            from tqdm.notebook import tqdm
+        except Exception:
+
+            def tqdm(x):
+                return x
+
+        for i in tqdm(range(xy.shape[0])):
+            if xy.shape[-1] == 2:
+                smoothed += (
+                    exp(-((norm((xy[i, :] - XY), axis=-1) / (sqrt(2) * sigma)) ** 2.0))
+                    / factor
+                )
+            elif xy.shape[-1] == 3:
+                smoothed += (
+                    exp(-((norm((xy[i, :2] - XY), axis=-1) / (sqrt(2) * sigma)) ** 2.0))
+                    / factor
+                    * xy[i, 2]
+                )
+    return smoothed, XY
+
+
+def plot_grid_data(xy_grid, data_grid, cbar=False, _plt=None, _cmap="RdBu", grid_min=1):
+    """# Plot Gridded Data"""
+    from numpy import nanmin, nanmax
+    from matplotlib.colors import TwoSlopeNorm
+
+    if _plt is None:
+        _plt = plt
+    norm = TwoSlopeNorm(
+        vmin=nanmin([-grid_min, nanmin(data_grid)]),
+        vcenter=0,
+        vmax=nanmax([grid_min, nanmax(data_grid)]),
+    )
+    _plt.pcolor(xy_grid[:, :, 0], xy_grid[:, :, 1], data_grid, cmap="RdBu", norm=norm)
+    if cbar:
+        _plt.colorbar()
+
+
+def plot_winding(xy_grid, winding_grid, cbar=False, _plt=None):
+    plot_grid_data(xy_grid, winding_grid, cbar=cbar, _plt=_plt, grid_min=1e-5)
+
+
+def plot_vorticity(xy_grid, vorticity_grid, cbar=False, save=False, _plt=None):
+    plot_grid_data(xy_grid, vorticity_grid, cbar=cbar, _plt=_plt, grid_min=1e-1)
+
+
+plt.blank = lambda: blank_plot(plt)
+plt.scatter2d = lambda data, labels=None, **matplotlib_kwargs: scatter_2d(
+    data, labels, plt, **matplotlib_kwargs
+)
+plt.save = save_plot
+
+
+def interpolate_points(df, columns=["Latent-0", "Latent-1"], kind="quadratic", N=500):
+    """# Interpolate points"""
+    from scipy.interpolate import interp1d
+    import numpy as np
+
+    points = df[columns].to_numpy()
+    distance = np.cumsum(np.sqrt(np.sum(np.diff(points, axis=0) ** 2, axis=1)))
+    distance = np.insert(distance, 0, 0) / distance[-1]
+    interpolator = interp1d(distance, points, kind=kind, axis=0)
+    alpha = np.linspace(0, 1, N)
+    interpolated_points = interpolator(alpha)
+    return interpolated_points.reshape(-1, 1, 2)
+
+
+def plot_trajectory(
+    df,
+    Fm=24,
+    FM=96,
+    interpolate=False,
+    interpolation_kind="quadratic",
+    interpolation_N=500,
+    columns=["Latent-0", "Latent-1"],
+    frame_column="Frames",
+    alpha=0.8,
+    lw=4,
+    _plt=None,
+    _z=None,
+):
+    """# Plot Trajectories
+
+    Interpolation possible"""
+    import numpy as np
+    from matplotlib.collections import LineCollection
+    import matplotlib as mpl
+
+    if _plt is None:
+        _plt = plt
+    if type(_plt) == mpl.axes._axes.Axes:
+        _ca = _plt
+    else:
+        try:
+            _ca = _plt.gca()
+        except:
+            _ca = _plt
+    X = df[columns[0]]
+    Y = df[columns[1]]
+    fm, fM = np.min(df[frame_column]), np.max(df[frame_column])
+
+    if interpolate:
+        if interpolation_kind == "cubic":
+            if len(df) <= 3:
+                return
+        elif interpolation_kind == "quadratic":
+            if len(df) <= 2:
+                return
+        else:
+            if len(df) <= 1:
+                return
+        points = interpolate_points(
+            df, kind=interpolation_kind, columns=columns, N=interpolation_N
+        )
+    else:
+        points = np.array([X, Y]).T.reshape(-1, 1, 2)
+
+    segments = np.concatenate([points[:-1], points[1:]], axis=1)
+    lc = LineCollection(
+        segments,
+        cmap=plt.get_cmap("plasma"),
+        norm=mpl.colors.Normalize(Fm, FM),
+    )
+    if _z is not None:
+        from mpl_toolkits.mplot3d.art3d import line_collection_2d_to_3d
+
+        if interpolate:
+            _z = _z  # TODO: Interpolate
+        line_collection_2d_to_3d(lc, _z)
+    lc.set_array(np.linspace(fm, fM, points.shape[0]))
+    lc.set_linewidth(lw)
+    lc.set_alpha(alpha)
+    _ca.add_collection(lc)
+    _ca.autoscale()
+    _ca.margins(0.04)
+    return lc
diff --git a/code/sunlab/common/plotting/colors.py b/code/sunlab/common/plotting/colors.py
new file mode 100644
index 0000000..c4fc727
--- /dev/null
+++ b/code/sunlab/common/plotting/colors.py
@@ -0,0 +1,38 @@
+class PhenotypeColors:
+    """# Phenotype Colorings
+
+    Standardization for the different phenotype colors"""
+
+    def __init__(self):
+        """# Empty Construtor"""
+        pass
+
+    def get_basic_colors(self, transition=False):
+        """# Return the Color Names
+
+        - transition: Returns the color for the transition class too"""
+        if transition:
+            return ["yellow", "purple", "green", "blue", "cyan"]
+        return ["yellow", "purple", "green", "blue"]
+
+    def get_colors(self, transition=False):
+        """# Return the Color Names
+
+        - transition: Returns the color for the transition class too"""
+        if transition:
+            return ["#ffff00", "#ff3cfa", "#11f309", "#213ff0", "cyan"]
+        return ["#ffff00", "#ff3cfa", "#11fe09", "#213ff0"]
+
+    def get_colormap(self, transition=False):
+        """# Return the Matplotlib Colormap
+
+        - transition: Returns the color for the transition class too"""
+        from matplotlib.colors import ListedColormap as LC
+
+        return LC(self.get_colors(transition))
+
+
+# Basic Exports
+Pcolor = PhenotypeColors().get_colors()
+Pmap = PhenotypeColors().get_colormap()
+Pmapx = PhenotypeColors().get_colormap(True)