From b85ee9d64a536937912544c7bbd5b98b635b7e8d Mon Sep 17 00:00:00 2001 From: Christian C Date: Mon, 11 Nov 2024 12:29:32 -0800 Subject: Initial commit --- code/sunlab/sunflow/plotting/model_extensions.py | 289 +++++++++++++++++++++++ 1 file changed, 289 insertions(+) create mode 100644 code/sunlab/sunflow/plotting/model_extensions.py (limited to 'code/sunlab/sunflow/plotting/model_extensions.py') 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 +) -- cgit v1.2.1