1 files changed, 289 insertions, 0 deletions
diff --git a/code/sunlab/sunflow/plotting/model_extensions.py b/code/sunlab/sunflow/plotting/model_extensions.py
new file mode 100644
index 0000000..087f8d3
--- /dev/null
+++ b/code/sunlab/sunflow/plotting/model_extensions.py
@@ -0,0 +1,289 @@
+from matplotlib import pyplot as plt
+from sunlab.common.data.shape_dataset import ShapeDataset
+from sunlab.globals import DIR_ROOT
+
+
+def get_nonphysical_masks(
+    model,
+    xrange=[-1, 1],
+    yrange=[-1, 1],
+    bins=[500, 500],
+    equivdiameter_threshold=10,
+    solidity_threshold=0.1,
+    area_threshold=100,
+    perimeter_threshold=10,
+    area_max_threshold=7000,
+    perimeter_max_threshold=350,
+    area_min_threshold=100,
+    perimeter_min_threshold=5,
+    consistency_check=False,
+):
+    """# Generate the Nonphysical Masks in Grid for Model
+
+    Hard Constraints:
+    - Non-negative values
+    - Ratios no greater than 1
+
+    Soft Constraints:
+    - Area/ Perimeter Thresholds"""
+    import numpy as np
+
+    x = np.linspace(xrange[0], xrange[1], bins[0])
+    y = np.linspace(yrange[0], yrange[1], bins[1])
+    X, Y = np.meshgrid(x, y)
+    X, Y = X.reshape((bins[0], bins[1], 1)), Y.reshape((bins[0], bins[1], 1))
+    XY = np.concatenate([X.reshape((-1, 1)), Y.reshape((-1, 1))], axis=-1)
+    dec_v = model.decoder(XY).numpy().reshape((bins[0] * bins[1], 13))
+    lXY = model.scaler.scaler.inverse_transform(dec_v).reshape((bins[0], bins[1], 13))
+    # Hard Limits
+    non_negative_mask = np.all(lXY > 0, axis=-1)
+    solidity_mask = np.abs(lXY[:, :, 6]) <= 1
+    extent_upper_bound_mask = lXY[:, :, 7] <= 1
+    # Soft Extremas
+    area_max_mask = lXY[:, :, 4] < area_max_threshold
+    perimeter_max_mask = lXY[:, :, 9] < perimeter_max_threshold
+    area_min_mask = lXY[:, :, 4] > area_min_threshold
+    perimeter_min_mask = lXY[:, :, 9] > perimeter_min_threshold
+    # Self-Consistency
+    man_solidity_mask = np.abs(lXY[:, :, 0] / lXY[:, :, 4]) <= 1
+    equivalent_diameter_mask = (
+        np.abs(lXY[:, :, 5] - np.sqrt(4 * np.abs(lXY[:, :, 0]) / np.pi))
+        < equivdiameter_threshold
+    )
+    convex_area_mask = lXY[:, :, 0] < lXY[:, :, 4] + area_threshold
+    convex_perimeter_mask = lXY[:, :, 9] < lXY[:, :, 8] + perimeter_threshold
+    mask_info = {
+        "non-negative": non_negative_mask,
+        "solidity": solidity_mask,
+        "extent-max": extent_upper_bound_mask,
+        #
+        "area-max": area_max_mask,
+        "perimeter-max": perimeter_max_mask,
+        "area-min": area_min_mask,
+        "perimeter-min": perimeter_min_mask,
+        #
+        "computed-solidity": man_solidity_mask,
+        "equivalent-diameter": equivalent_diameter_mask,
+        "convex-area": convex_area_mask,
+        "convex-perimeter": convex_perimeter_mask,
+    }
+    if not consistency_check:
+        mask_info = {
+            "non-negative": non_negative_mask,
+            "solidity": solidity_mask,
+            "extent-max": extent_upper_bound_mask,
+            #
+            "area-max": area_max_mask,
+            "perimeter-max": perimeter_max_mask,
+            "area-min": area_min_mask,
+            "perimeter-min": perimeter_min_mask,
+        }
+    mask_list = [mask_info[key] for key in mask_info.keys()]
+    return np.all(mask_list, axis=0), X, Y, mask_info
+
+
+def excavate(input_2d_array):
+    """# Return Boundaries for Masked Array
+
+    Use X, Y directions only"""
+    from copy import deepcopy as dc
+    from numpy import nan_to_num, zeros_like, abs
+
+    data_2d_array = dc(input_2d_array)
+    data_2d_array = nan_to_num(data_2d_array, nan=20)
+    # X-Gradient
+    x_grad = zeros_like(data_2d_array)
+    x_grad[:-1, :] = data_2d_array[1:, :] - data_2d_array[:-1, :]
+    x_grad[(abs(x_grad) > 10)] = 10
+    x_grad[(abs(x_grad) < 10) & (abs(x_grad) > 0)] = 1
+    x_grad[x_grad == 1] = 0.5
+    x_grad[x_grad > 1] = 1
+    # Y-Gradient
+    y_grad = zeros_like(data_2d_array)
+    y_grad[:, :-1] = data_2d_array[:, 1:] - data_2d_array[:, :-1]
+    y_grad[(abs(y_grad) > 10)] = 10
+    y_grad[(abs(y_grad) < 10) & (abs(y_grad) > 0)] = 1
+    y_grad[y_grad == 1] = 0.5
+    y_grad[y_grad > 1] = 1
+    return x_grad, y_grad
+
+
+def excavate_extra(input_2d_array, N=1):
+    """# Return Boundaries for Masked Array
+
+    Use all 8 directions"""
+    from copy import deepcopy as dc
+    from numpy import nan_to_num, zeros_like, abs
+
+    data_2d_array = dc(input_2d_array)
+    data_2d_array = nan_to_num(data_2d_array, nan=20)
+    # X-Gradient
+    x_grad = zeros_like(data_2d_array)
+    x_grad[:-N, :] = data_2d_array[N:, :] - data_2d_array[:-N, :]
+    x_grad[(abs(x_grad) > 10)] = 10
+    x_grad[(abs(x_grad) < 10) & (abs(x_grad) > 0)] = 1
+    x_grad[x_grad == 1] = 0.5
+    x_grad[x_grad > 1] = 1
+    # Y-Gradient
+    y_grad = zeros_like(data_2d_array)
+    y_grad[:, :-N] = data_2d_array[:, N:] - data_2d_array[:, :-N]
+    y_grad[(abs(y_grad) > 10)] = 10
+    y_grad[(abs(y_grad) < 10) & (abs(y_grad) > 0)] = 1
+    y_grad[y_grad == 1] = 0.5
+    y_grad[y_grad > 1] = 1
+    # XY-Gradient
+    xy_grad = zeros_like(data_2d_array)
+    xy_grad[:-N, :-N] = data_2d_array[N:, N:] - data_2d_array[:-N, :-N]
+    xy_grad[(abs(xy_grad) > 10)] = 10
+    xy_grad[(abs(xy_grad) < 10) & (abs(xy_grad) > 0)] = 1
+    xy_grad[xy_grad == 1] = 0.5
+    xy_grad[xy_grad > 1] = 1
+    # X(-Y)-Gradient
+    yx_grad = zeros_like(data_2d_array)
+    yx_grad[:-N, :-N] = data_2d_array[N:, :-N] - data_2d_array[:-N, N:]
+    yx_grad[(abs(yx_grad) > 10)] = 10
+    yx_grad[(abs(yx_grad) < 10) & (abs(yx_grad) > 0)] = 1
+    yx_grad[yx_grad == 1] = 0.5
+    yx_grad[yx_grad > 1] = 1
+    xyn_grad = dc(yx_grad)
+    # (-X)Y-Gradient
+    xny_grad = zeros_like(data_2d_array)
+    xny_grad[:-N, :-N] = data_2d_array[:-N, N:] - data_2d_array[N:, :-N]
+    xny_grad[(abs(xy_grad) > 10)] = 10
+    xny_grad[(abs(xy_grad) < 10) & (abs(xy_grad) > 0)] = 1
+    xny_grad[xy_grad == 1] = 0.5
+    xny_grad[xy_grad > 1] = 1
+    # (-X)(-Y)-Gradient
+    xnyn_grad = zeros_like(data_2d_array)
+    xnyn_grad[:-N, :-N] = data_2d_array[:-N, :-N] - data_2d_array[N:, N:]
+    xnyn_grad[(abs(yx_grad) > 10)] = 10
+    xnyn_grad[(abs(yx_grad) < 10) & (abs(yx_grad) > 0)] = 1
+    xnyn_grad[yx_grad == 1] = 0.5
+    xnyn_grad[yx_grad > 1] = 1
+    return x_grad, y_grad, xy_grad, xyn_grad, xny_grad, xnyn_grad
+
+
+def excavate_outline(arr, thickness=1):
+    """# Generate Transparency Mask with NaNs"""
+    from numpy import sum, abs, NaN
+
+    outline = sum(abs(excavate_extra(arr, thickness)), axis=0)
+    outline[outline == 0] = NaN
+    outline[outline > 0] = 0
+    return outline
+
+
+def get_boundary_outline(
+    aae_model_object,
+    pixel_classification_file=None,
+    include_transition_regions=False,
+    border_thickness=3,
+    bin_count=800,
+    xrange=[-6.5, 6.5],
+    yrange=[-4.5, 4.5],
+    threshold=0.75,
+):
+    """# Get Boundary Outlines"""
+    from copy import deepcopy
+    import numpy as np
+
+    if pixel_classification_file is None:
+        pixel_classification_file = "../../extra_data/PhenotypePixels_65x45_800.npy"
+    base_classification = np.loadtxt(pixel_classification_file)
+    base_classification = base_classification.reshape((bin_count, bin_count, 4))
+    max_classification_probability = np.zeros((bin_count, bin_count, 1))
+    max_classification_probability[:, :, 0] = (
+        np.max(base_classification, axis=-1) < threshold
+    )
+    classes_with_include_transition_regions = np.concatenate(
+        [base_classification, max_classification_probability], axis=-1
+    )
+    if include_transition_regions:
+        phenotype_probabilities = deepcopy(
+            np.argsort(classes_with_include_transition_regions[:, :, :], axis=-1)[
+                :, :, -1
+            ]
+        ).astype(np.float32)
+    else:
+        phenotype_probabilities = deepcopy(
+            np.argsort(classes_with_include_transition_regions[:, :, :-1], axis=-1)[
+                :, :, -1
+            ]
+        ).astype(np.float32)
+    nonphysical_mask, _, _, _ = get_nonphysical_masks(
+        aae_model_object, xrange=xrange, yrange=yrange, bins=[bin_count, bin_count]
+    )
+    nonphysical_mask = nonphysical_mask.astype(np.float32)
+    nonphysical_mask[nonphysical_mask == 0] = np.NaN
+    nonphysical_mask[nonphysical_mask == 1] = 0
+    nonphysical_mask = nonphysical_mask.T
+    phenotype_regions = deepcopy(phenotype_probabilities.T + nonphysical_mask.T)
+    outline = excavate_outline(phenotype_regions, border_thickness)
+    return outline
+
+
+def apply_boundary(
+    model_loc=DIR_ROOT + "models/current_model/",
+    border_thickness=3,
+    include_transition_regions=False,
+    threshold=0.7,
+    alpha=1,
+    _plt=None,
+):
+    """# Apply Boundary to Plot
+
+    Use Pregenerated Boundary by Default for Speed"""
+    from ..models import load_aae
+    from sunlab.common.scaler import MaxAbsScaler
+    import numpy as np
+
+    if _plt is None:
+        _plt = plt
+    if (model_loc == model_loc) and (border_thickness == 3) and (threshold == 0.7):
+        XYM = np.load(DIR_ROOT + "extra_data/OutlineXYM.npy")
+        XY = XYM[:2, :, :]
+        if include_transition_regions:
+            outline = XYM[3, :, :]
+        else:
+            outline = XYM[2, :, :]
+        _plt.pcolor(XY[0, :, :], XY[1, :, :], outline, cmap="gray", alpha=alpha)
+        return
+    model = load_aae(model_loc, MaxAbsScaler)
+    bin_count = 800
+    xrange = [-6.5, 6.5]
+    yrange = [-4.5, 4.5]
+    rng = [xrange, yrange]
+    X = np.linspace(rng[0][0], rng[0][1], bin_count)
+    Y = np.linspace(rng[1][0], rng[1][1], bin_count)
+    XY = np.array(np.meshgrid(X, Y))
+    kwparams = {
+        "bin_count": bin_count,
+        "xrange": xrange,
+        "yrange": yrange,
+    }
+
+    include_tregions = include_transition_regions
+    outline = get_boundary_outline(
+        model,
+        border_thickness=border_thickness,
+        include_transition_regions=include_tregions,
+        threshold=threshold,
+        **kwparams
+    )
+    _plt.pcolor(XY[0, :, :], XY[1, :, :], outline, cmap="gray", alpha=alpha)
+
+
+plt.apply_boundary = apply_boundary
+
+
+def plot_shape_dataset(self, model, *args, **kwargs):
+    """# Plot Shape Dataset"""
+    if self.labels is None:
+        plt.scatter2d(model.encoder(self.dataset), *args, **kwargs)
+    else:
+        plt.scatter2d(model.encoder(self.dataset), self.labels, *args, **kwargs)
+
+
+ShapeDataset.plot = lambda model, *args, **kwargs: plot_shape_dataset(
+    model, *args, **kwargs
+)