Работа с известными индексами в массиве Numpy

Примечание. Я свел объяснения к минимуму, но этот ответ все еще длинный. Если вы ищете только код, просто прокрутите до конца.

По моему мнению, ключом к векторизации является извлечение частей, которые можно векторизовать.

Но перед этим есть одна часть, которую необходимо исправить. Следующий код делает что-то очень избыточное.

values, counts = np.unique_counts(self.neighbours(index, 1))
probability = counts.astype(np.float16)
probability /= np.sum(probability)
random = self.rng.choice(values, p=probability)

Его можно заменить следующим:

random = self.rng.choice(self.neighbours(index, 1).ravel())

Например, если соседей 9 и 3 из них — вода, вероятность выбора воды будет 3/9. Однако обратите внимание, что результат не будет таким же, как оригинал, хотя математически он делает то же самое.

Теперь, поскольку операция, которую вы хотели выполнить, заключалась в случайном выборе одного из соседей, вы можете напрямую получить доступ к целевой плитке следующим образом:

neighbour_offset = np.array([(dx, dy) for dx in (-1, 0, 1) for dy in (-1, 0, 1)])
neighbour_index = self.rng.choice(neighbour_offset) + index
random = self.world_map[tuple(neighbour_index)]

Вы можете видеть, что rng.choice нужно выбрать только один из neighbour_offset, который не меняется во время цикла. Следовательно, это можно вычислить вне цикла. Ниже приведены необходимые изменения:

    def check_tile(self, index: tuple, neighbour_indexes):
        """
        Texturize borders and update new borders.

        index: tuple representing the index of the tile to check.
        """
        # Corresponding land id to current working tile.
        tile = self.world_map[index]
        # If tile at the index is surrounded by identical tiles within a 2-tile range.
        if np.all(self.neighbours(index, 2) == tile):
            # Don't store it; this tile cannot become water because it is surrounded by 2-tile wide same land as itself.
            pass
        else:
            # Randomly change each of the 4 newly descended tiles of the original tile to either water or not.
            for row in range(2):
                for column in range(2):
                    # One of the four descended tile's index.
                    new_tile_index = (index[0] * 2 + row, index[1] * 2 + column)
                    # Randomly chose a value according to its probability.
                    neighbour_index = neighbour_indexes[row * 2 + column]
                    random = self.world_map[tuple(neighbour_index)]
                    if random == 0:  # If tile at the index became water.
                        # Update it on the new map.
                        self.newmap[new_tile_index] = random
                        # Don't store it because it is a water tile and no longer a border tile.
                    elif random == tile:  # If the tile remained the same.
                        # Store it because it is still a border tile.
                        self.newborders.append(new_tile_index)
                    else:  # If the tile changed to a different land.
                        # Update it on the new map.
                        self.newmap[new_tile_index] = random
                        # Store it because it is still a border tile.
                        self.newborders.append(new_tile_index)

    def default_procedure(self):
        """
        Default procedure: upscale (or zoom into) the map and texturize borders.

        """
        self.upscale_map()

        # Choose 4 random neighbours for each border tile.
        # Note that below is NOT taking into account the case where the land is at the edge of the map.
        neighbour_offset = np.array([(dx, dy) for dx in (-1, 0, 1) for dy in (-1, 0, 1)])
        random_neighbour_offsets = self.rng.choice(neighbour_offset, size=(len(self.borders), 4))
        random_neighbour_indexes = np.asarray(self.borders)[:, None, :] + random_neighbour_offsets

        for index, neighbour_index in zip(self.borders, random_neighbour_indexes):
            self.check_tile(index, neighbour_index)
        self.borders = self.newborders
        self.newborders = []
        self.world_map = self.newmap

(Обратите внимание, что в этом коде не учитываются тайлы земли на краю карты. Я опустил здесь этот случай, потому что это предотвратит векторизацию. Если вам нужно с этим справиться, я рекомендую разделить процесс на несколько шагов. и обрабатываем края отдельно.)

Благодаря этому rng.choice теперь векторизован. На моем ПК только этот шаг сделал его в 8 раз быстрее, чем оригинал. Вы можете продолжать этот процесс неоднократно. Наверное, мне следует все объяснить, но это будет очень долго, поэтому, пожалуйста, позвольте мне пропустить до конца.

Это полностью векторизованный код:

    def default_procedure(self):
        """
        Default procedure: upscale (or zoom into) the map and texturize borders.

        """
        self.upscale_map()

        def take(arr, indexes):
            return arr[indexes[..., 0], indexes[..., 1]]

        def put(arr, indexes, values):
            arr[indexes[..., 0], indexes[..., 1]] = values

        land_tile_indexes = np.asanyarray(self.borders)
        land_tiles = take(self.world_map, land_tile_indexes)

        # Skip tiles that are surrounded by water.
        neighbour_25_windows = np.lib.stride_tricks.sliding_window_view(
            np.pad(self.world_map, pad_width=5 // 2, mode='edge'),  # Pad the map in order to skip edge cases.
            window_shape=(5, 5),
        )
        neighbour_25_windows = take(neighbour_25_windows, land_tile_indexes)  # Filter only the border tiles.
        is_not_surrounded_tile = np.any(neighbour_25_windows != land_tiles[..., None, None], axis=(-2, -1))

        # Filter land_tile_indexes by the condition.
        land_tile_indexes = land_tile_indexes[is_not_surrounded_tile]

        # Choose 4 random neighbours for each border tile.
        # Note that below is NOT taking into account the case where the land is at the edge of the map.
        neighbour_offset = np.array([(dx, dy) for dx in (-1, 0, 1) for dy in (-1, 0, 1)])
        random_neighbour_offsets = self.rng.choice(neighbour_offset, size=(len(land_tile_indexes), 4))
        random_neighbour_indexes = land_tile_indexes[:, None, :] + random_neighbour_offsets
        random_neighbour_tiles = take(self.world_map, random_neighbour_indexes)

        # Pre-calculate all the new tile indexes here.
        upscaling_offset = np.array([(dx, dy) for dx in range(2) for dy in range(2)])
        new_tile_indexes = (land_tile_indexes * 2)[:, None, :] + upscaling_offset[None]

        # Put the new tile if it did NOT remain the same.
        # In other words, this covers below cases:
        #  - If the tile became water.
        #  - If the tile changed to a different land.
        condition = random_neighbour_tiles != take(self.world_map, land_tile_indexes)[..., None]
        put(self.newmap, new_tile_indexes[condition], random_neighbour_tiles[condition])

        # Update the border tiles if it did NOT become water.
        # In other words, this covers below cases:
        #  - If the tile remained the same.
        #  - If the tile changed to a different land.
        target_indexes = new_tile_indexes[random_neighbour_tiles != 0]

        self.borders = target_indexes
        self.world_map = self.newmap

На данный момент узкое место self.upscale_map(). На моем ПК этот процесс занял более 80% времени выполнения. Самый простой способ решить эту проблему — использовать OpenCV.

    def upscale_map(self) -> None:
        """
        Divide each tile into 4 pieces and upscale the map by a factor of 2.

        """
        self.newmap = cv2.resize(self.world_map, (0, 0), fx=2, fy=2, interpolation=cv2.INTER_NEAREST)

Вот результат после 13 итераций.

Наконец, в качестве альтернативы вы можете использовать Нумба. С Numba реализация может оказаться более интуитивно понятной. Однако объединение Numba с классами может быть немного сложным, поэтому вам все равно придется переписать большую часть кода. Тем не менее, на это стоит обратить внимание.

Я не думаю, что этот код точно воспроизводит ваши правила, но я думаю, что идея реализации этого с помощью высокоуровневых операций масштабирования, фильтрации и бинаризации может помочь вам добиться большей производительности, чем узкие for циклы.

import numpy as np
from scipy import ndimage
import matplotlib.pyplot as plt
rng = np.random.default_rng(1)

A = np.zeros((4, 4)).astype(bool)
A[1:3, 1:3] = 1

for j in range(8):
    A = ndimage.zoom(A, 2, order=0)
    B = ndimage.uniform_filter(A.astype(float))
    z = rng.random(A.sum())
    A[A] = B[A] > z

plt.imshow(A)

Это выполняет массу избыточной работы, но я думаю, что это все равно немного быстрее, чем у вас. Если добавить немного дополнительной логики, чтобы избежать фильтрации больших блоков, состоящих из воды, это, вероятно, может быть довольно быстрым, особенно если вы можете использовать cupyx.scipy.ndimage с графическим процессором.