Numpy APIs

def extract_3d_shapes_from_model_file(
    input_file: str,
    layer_names: list,
    max_layers: int = None,
    device: str = "cpu",
    target_faces: int = None,
    target_rows: int = 512,
    target_cols: int = 512,
    start_x: int = 0,
    start_y: int = 0,
    start_z: int = 1,
    max_depth: int = 2,
    pad_per: int = 20,
):
    """
    extract_3d_shapes_from_model_file

    extract the weights from a model file
    and build a list of 3d arrays with
    an a marching cubes algorithm that
    finds shapes in high resolution data

    :param input_file: path to model.safetensors
        file
    :param layer_names: filter by layer
        layer colomn names
    :param max_layers: limit the number of
        layers to fit
    :param device: cpu vs gpu
    :param target_faces: min faces to
        hopefully fit if there is
        enough data
    :param target_rows: number of
        resample target rows before
        drawing
    :param target_cols: number of
        resample target cols before
        drawing
    :param max_depth: stack each layer
        on itself this many times
        to convert it to a 3d ndarray
    :param pad_per: number to pad
        per object
        on the y-axis
    """
    tensors = {}
    tensor_keys = []
    num_tensors = 0

    log.info(
        f"extracting tensors from model={input_file} "
        f"layers={','.join(tensor_keys)}"
    )
    tensors = get_model_tensors.get_model_tensors(
        input_file, layer_names, device=device
    )

    """
    print("Model Tensors:")
    print(tensors)
    print(f"num tensors: {len(tensors)}")
    """
    num_tensors = len(tensors)
    log.info(
        "preprocess - phase 1 - "
        f"reading {num_tensors} tensors "
        f"filtering layers={len(layer_names)}"
    )

    target_size_mb = float(
        f"{float(target_rows * target_cols * 4.0 / 1024.0 / 102.4)}"
    )

    # set some depth for shape rendering with smaller renders
    if max_layers is not None and max_layers < 5:
        max_depth = 3

    all_data_3d = []
    tensor_2d_arrays = []
    num_fit = 0
    for idx, key in enumerate(tensors):
        tensor_node = tensors[key]
        # https://github.com/pytorch/pytorch/issues/110285
        try:
            tensor_data = tensor_node["data"].numpy()
        except Exception as e:
            log.debug(
                f"ignored tensor={idx} {key} with ex={e}"
            )
            continue
        # refresh to make this faster
        tensor_2d_arrays = []
        tensor_2d_arrays.append(tensor_data)
        tensor_2d_arrays.append(tensor_data)
        tensor_num_rows = tensor_data.shape[0]
        tensor_num_cols = tensor_data.shape[1]
        for i in range(0, max_depth):
            stacked_arr = np.tile(
                tensor_data[:, :, np.newaxis], (1, 1, 3)
            )
        if len(stacked_arr) > 1:
            mb_size_org = float(stacked_arr.nbytes) / (
                1024.0 * 1024.0
            )
            layer_name = f"{key[0:128]}"
            label_layer_name = f"Layer: {layer_name}"
            desc = (
                f"src dimensions=({tensor_num_rows}, "
                f"{tensor_num_cols}) "
                f"size={mb_size_org:.2f}mb compressed to "
                f"({target_rows},{target_cols}) "
                f"{target_size_mb:.2f}mb"
            )
            if (num_tensors < 50) or (idx % 500 == 0):
                log.info(
                    f"fitting {num_fit + 1}/{num_tensors} "
                    f"{desc} dst=({target_rows}, {target_cols}) "
                    f"{target_size_mb:.2f}mb"
                    ""
                )
            fitted_3d_array = (
                fit.fit_2d_arrays_to_target_shape(
                    tensor_2d_arrays,
                    target_rows,
                    target_cols,
                )
            )
            all_data_3d.append(
                {
                    "name": layer_name,
                    "layer_name": label_layer_name,
                    "desc": desc,
                    "data": fitted_3d_array,
                    "target_faces": target_faces,
                    "target_rows": target_rows,
                    "target_cols": target_cols,
                    "x": start_x,
                    "y": start_y + (num_fit * pad_per),
                    "z": start_z,
                }
            )
            num_fit += 1
            if max_layers:
                if num_fit >= max_layers:
                    break
    # for key in tensors

    if max_layers:
        log.info(
            f"done fitting {num_fit}/{max_layers} "
            f"out of {num_tensors} "
            f"into ~{target_faces} polygon faces "
            "per tensor in "
            f"shape({target_rows}, {target_cols})"
        )
    else:
        log.info(
            f"done fitting {num_fit}/{num_tensors} "
            f"into ~{target_faces} polygon faces "
            "per tensor in "
            f"shape({target_rows}, {target_cols})"
        )
    return all_data_3d

def fit_2d_arrays_to_target_shape(
    arrays: list[np.ndarray],
    target_rows: int,
    target_cols: int,
):
    """
    fit_2d_arrays_to_target_shape

    fit a list of 2D arrays into a target 3D array with specified dimensions.

    :param arrays: List of 2D NumPy arrays.
    :param target_rows: number of
        resample target rows before
        drawing
    :param target_cols: number of
        resample target cols before
        drawing

    :return: 3d np.ndarray with the resized 2d arrays stacked on
        the z-axis
    """
    num_arrays = len(arrays)
    fitted_arrays = []
    for idx, array in enumerate(arrays):
        op_performed = None
        if (
            array.shape[0] < target_rows
            or array.shape[1] < target_cols
        ):
            op_performed = "upscaled"
            log.debug(
                f"fit {idx}/{num_arrays} "
                f"upscaled {array.shape}"
            )
            fitted_array = upscaler.upscale_2d_array(
                array, target_rows, target_cols
            )
        elif (
            array.shape[0] > target_rows
            or array.shape[1] > target_cols
        ):
            op_performed = "downscale"
            log.debug(
                f"fit {idx}/{num_arrays} "
                f"downscaled {array.shape}"
            )
            fitted_array = downscaler.downscale_2d_array(
                array, target_rows, target_cols
            )
        else:
            op_performed = "ignored"
            log.debug(
                f"fit {idx}/{num_arrays} "
                f"ignored {array.shape}"
            )
            fitted_array = array

        log.debug(
            f"fit {idx}/{num_arrays} "
            f"{op_performed} "
            f"src={array.shape} "
            f"dst={fitted_array.shape} == "
            f"({target_rows}, {target_cols})"
        )
        fitted_arrays.append(fitted_array)

    return np.stack(fitted_arrays, axis=2)

def downscale_2d_array(
    array: np.ndarray, target_rows: int, target_cols: int
):
    """
    downscale_2d_array

    downscale a 2D array to the target dimensions by averaging values

    :param array: input 2D NumPy array
    :param target_rows: number of rows
    :param target_cols: number of columns

    :return: downscaled 2D numpy array
    """
    downscale_factor_rows = max(
        1, array.shape[0] // target_rows
    )
    downscale_factor_cols = max(
        1, array.shape[1] // target_cols
    )

    reshaped_array = array[
        : downscale_factor_rows * target_rows,
        : downscale_factor_cols * target_cols,
    ]
    log.debug(
        f"downscaling {array.shape} "
        f"factor_rows={downscale_factor_rows} "
        f"factor_cols={downscale_factor_cols} "
    )
    reshaped_array = reshaped_array.reshape(
        target_rows,
        downscale_factor_rows,
        target_cols,
        downscale_factor_cols,
    )

    downscaled_array = np.mean(reshaped_array, axis=(1, 3))

    return downscaled_array

def upscale_2d_array(
    array: np.ndarray,
    target_rows: int,
    target_cols: int,
):
    """
    upscale_2d_array

    upscale a 2D array to the target dimensions by repeating values

    :param array: input 2D NumPy array
    :param target_rows: number of rows
    :param target_cols: number of columns

    :return: upscaled 2D numpy array
    """
    scale_factor_rows = target_rows / array.shape[0]
    scale_factor_cols = target_cols / array.shape[1]
    # use np.tile to repeat the array along each dimension
    upscaled_array = np.tile(
        array,
        (
            int(np.ceil(scale_factor_rows)),
            int(np.ceil(scale_factor_cols)),
        ),
    )

    # crop the upscaled array to match the target dimensions
    upscaled_array = upscaled_array[
        :target_rows, :target_cols
    ]
    return upscaled_array

def calculate_weighted_quantile_ranges_2d(
    array_2d: np.ndarray,
    num_ranges: int,
):
    """
    calculate_weighted_quantile_ranges_2d

    build a quantile range based off the 2d array
    z-axis values and the number of ranges. assign
    colors based off an ordered list from
    bw.bl.colors.get_color_tuples()

    :param array_2d: 2d array data to process
    :param num_ranges: number of np.quantile ranges
        to calculate min/max lower/upper bounds
        for quickly assigning colors to a z-value
        in the array_2d
    """
    # Calculate the minimum and maximum values for all elements in the array
    min_value = int(np.min(array_2d))
    max_value = int(np.max(array_2d)) + 1

    # Calculate quantile ranges for array values
    quantile_ranges = np.linspace(
        min_value, max_value, num_ranges + 1
    )

    # Create a dictionary to store colors for each quantile range
    colors_dict = {}

    # Apply seaborn color gradient with opacity 1.0
    # cmap = sns.color_palette("viridis", num_ranges)
    # cmap = sns.color_palette("bright", num_ranges)
    # colors = sns.color_palette(cmap, n_colors=num_ranges)

    colors_tuples_map = bwcl.get_color_tuples()
    colors_names = [
        color_name for color_name in colors_tuples_map
    ]
    colors_list = [
        colors_tuples_map[color_name]
        for color_name in colors_names
    ]

    # Assign colors to the dictionary based on quantile ranges
    num_colors = len(colors_list)
    cidx = 0
    for i in range(num_ranges):
        min_range_org = quantile_ranges[i]
        max_range_org = quantile_ranges[i + 1]
        min_range = float(f"{min_range_org:.2f}")
        max_range = float(f"{max_range_org:.2f}")
        color_key = f"{min_range}_{max_range}"
        color_node = colors_list[cidx]
        color_name = colors_names[cidx]
        """
        log.debug(
            f'{i}/{num_ranges} quantile '
            f'quantile=[{min_range},{max_range}]'
            f'color={color_name} {color_key}')
        """
        colors_dict[color_key] = {
            "name": colors_names[cidx],
            "r": color_node[0],
            "b": color_node[1],
            "g": color_node[2],
            "a": color_node[3],
            "q": i,
            "min": min_range,
            "max": max_range,
        }
        cidx += 1
        if cidx >= num_colors:
            cidx = 0

    for idx, color_range in enumerate(colors_dict):
        color_node = colors_dict[color_range]
        color_name = color_node["name"]
        min_range_val = color_node["min"]
        max_range_val = color_node["max"]
        log.debug(
            f"{idx} = [{min_range_val}, {max_range_val}] "
            f"{color_name} total=[{min_value}, {max_value}]"
        )
    return colors_dict

def calculate_weighted_quantile_ranges_3d(
    array_3d: np.ndarray,
    num_ranges: int,
):
    """
    calculate_weighted_quantile_ranges_3d

    build a quantile range based off the 3d array
    z-axis values and the number of ranges. assign
    colors based off an ordered list from
    bw.bl.colors.get_color_tuples()

    :param array_3d: 3d array data to process
    :param num_ranges: number of np.quantile ranges
        to calculate min/max lower/upper bounds
        for quickly assigning colors to a z-value
        in the array_3d
    """
    # Calculate the minimum and maximum values for all z values
    min_z = np.min(array_3d[:, :, 2])
    max_z = np.max(array_3d[:, :, 2])

    # Calculate quantile ranges for z values
    quantile_ranges = np.linspace(
        min_z, max_z, num_ranges + 1
    )

    # Create a dictionary to store colors for each quantile range
    colors_dict = {}

    # Apply seaborn color gradient with opacity 1.0
    # cmap = sns.color_palette("viridis", num_ranges)
    # cmap = sns.color_palette("bright", num_ranges)
    # colors = sns.color_palette(cmap, n_colors=num_ranges)

    colors_tuples_map = bwcl.get_color_tuples()
    colors_names = [
        color_name for color_name in colors_tuples_map
    ]
    colors_list = [
        colors_tuples_map[color_name]
        for color_name in colors_names
    ]

    # Assign colors to the dictionary based on quantile ranges
    num_colors = len(colors_list)
    cidx = 0
    for i in range(num_ranges):
        min_range_org = quantile_ranges[i]
        max_range_org = quantile_ranges[i + 1]
        min_range = float(f"{min_range_org:.2f}")
        max_range = float(f"{max_range_org:.2f}")
        color_key = f"{min_range}_{max_range}"
        color_node = colors_list[cidx]
        color_name = colors_names[cidx]
        """
        log.debug(
            f'{i}/{num_ranges} '
            f'quantile=[{min_range},{max_range}]'
            f'color={color_name} {color_key}')
        """
        colors_dict[color_key] = {
            "name": colors_names[cidx],
            "r": color_node[0],
            "b": color_node[1],
            "g": color_node[2],
            "a": color_node[3],
            "q": i,
            "min": min_range,
            "max": max_range,
        }
        cidx += 1
        if cidx >= num_colors:
            cidx = 0

    for idx, color_range in enumerate(colors_dict):
        color_node = colors_dict[color_range]
        color_name = color_node["name"]
        min_range_val = color_node["min"]
        max_range_val = color_node["max"]
        log.debug(
            f"{idx} = [{min_range_val}, {max_range_val}] "
            f"{color_name}"
        )
    return colors_dict

def create_random_2d_arrays(
    num_arrays: int,
    min_rows: int,
    max_rows: int,
    min_cols: int,
    max_cols: int,
):
    """
    create testing 2d arrays

    Create x number of random-sized 2D arrays.

    :param num_arrays: number of 2d arrays
        to create
    :param min_rows: min rows
    :param max_rows: max rows
    :param min_cols: min columns
    :param max_cols: max columns

    :return: list of 2d numpy ndarrays of float32 data
    """
    arrays = []
    for i in range(num_arrays):
        rows = np.random.randint(min_rows, max_rows + 1)
        cols = np.random.randint(min_cols, max_cols + 1)
        array = np.random.rand(rows, cols).astype(
            np.float32
        )
        mb_size = (
            float(array.shape[0] * array.shape[1] * 4)
            / 1024.0
            / 1024.0
        )
        log.info(
            f"created {i + 1}/{num_arrays} {array.shape} "
            f"{mb_size:.2f}mb"
        )
        arrays.append(array)
    return arrays