def square_log_axes(ax, nx, ny):
    """
    Transform log-scaled axes into a square aspect ratio.
    The axes will be resized so that each cell in the heatmap is square.
    TODO: This will work, but only if tight_layout() is not called
    """
    ratio = ny / nx
    pos1 = ax.get_position()
    fig_width = ax.get_figure().get_figwidth()
    fig_height = ax.get_figure().get_figheight()
    axratio = fig_height * pos1.height / (fig_width * pos1.width)
    pos2 = [pos1.x0, pos1.y0, pos1.width, pos1.height]
    if axratio > ratio:
        target_height = ratio * fig_width * pos1.width / fig_height
        delta_height = pos1.height - target_height
        pos2[1] = pos2[1] - delta_height / 2.0
        pos2[3] = target_height
        ax.set_position(pos2)
    elif axratio < ratio:
        print('Decreasing width')
        target_width = fig_height * pos1.height / (fig_width * ratio)
        delta_width = pos1.width - target_width
        pos2[0] = pos2[0] + delta_width / 2.0
        pos2[2] = target_width
        ax.set_position(pos2)
        print(ratio, axratio, (pos1.x0, pos1.y0, pos1.width, pos1.height), pos2
            )
