ml/vae/codec.py - main - sususer/ColonyGEN

ml/vae/codec.py
261 строка10 KB
Chekr
f11
03 июн 2026, 13:53
03 июн 2026, 13:53d260831
from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np

@dataclass
class CodecConfig:
    map_width: int = 48
    map_height: int = 48
    max_resource_value: int = 7

    def to_dict(self) -> dict[str, Any]:
        return {
            "map_width": self.map_width,
            "map_height": self.map_height,
            "max_resource_value": self.max_resource_value,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "CodecConfig":
        return cls(
            map_width=int(data.get("map_width", 48)),
            map_height=int(data.get("map_height", 48)),
            max_resource_value=int(data.get("max_resource_value", 7)),
        )

class WorldTensorCodec:
    terrain_types = (
        "base",
        "beach",
        "forest",
        "mountain",
        "plain",
        "sand",
        "water",
    )
    biome_types = (
        "coast",
        "colony_core",
        "dense_forest",
        "desert",
        "grassland",
        "highland",
        "lake",
        "mountain_range",
        "river",
        "wetland",
        "young_forest",
    )
    layer_names = (
        "elevation",
        "moisture",
        "fertility",
        "hazard",
        "influence",
    )
    terrain_resource_caps = {
        "forest": 2,
        "mountain": 5,
    }

    def __init__(self, config: CodecConfig | None = None):
        self.config = config or CodecConfig()
        self.terrain_to_idx = {name: idx for idx, name in enumerate(self.terrain_types)}
        self.biome_to_idx = {name: idx for idx, name in enumerate(self.biome_types)}
        self.num_tiles = self.config.map_width * self.config.map_height
        self.terrain_slice = slice(0, len(self.terrain_types))
        self.biome_slice = slice(self.terrain_slice.stop, self.terrain_slice.stop + len(self.biome_types))
        self.resource_index = self.biome_slice.stop
        self.base_index = self.resource_index + 1
        self.layer_slice = slice(self.base_index + 1, self.base_index + 1 + len(self.layer_names))
        self.feature_dim = self.layer_slice.stop
        self.feature_weights = np.ones(self.feature_dim, dtype=np.float32)
        self.feature_weights[self.terrain_slice] = 4.0
        self.feature_weights[self.terrain_slice.start + self.terrain_to_idx["water"]] = 6.2
        self.feature_weights[self.biome_slice] = 2.5
        self.feature_weights[self.resource_index] = 2.5
        self.feature_weights[self.base_index] = 3.0
        self.feature_weights[self.layer_slice] = 1.0

    @classmethod
    def from_manifest(cls, manifest: dict[str, Any]) -> "WorldTensorCodec":
        config = manifest.get("config", {})
        return cls(
            CodecConfig(
                map_width=int(config.get("map_width", 48)),
                map_height=int(config.get("map_height", 48)),
                max_resource_value=int(config.get("max_resource_value", 7)),
            )
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "config": self.config.to_dict(),
            "terrain_types": list(self.terrain_types),
            "biome_types": list(self.biome_types),
            "layer_names": list(self.layer_names),
        }

    def resource_cap_for_tile(self, terrain: str, biome: str | None = None) -> int:
        del biome
        return int(self.terrain_resource_caps.get(str(terrain).lower(), 0))

    def encode_resource_value(self, terrain: str, biome: str | None, resource: Any) -> float:
        cap = self.resource_cap_for_tile(terrain, biome)
        if cap <= 0:
            return 0.0
        return float(np.clip(float(resource) / float(cap), 0.0, 1.0))

    def decode_resource_value(self, terrain: str, biome: str | None, resource_value: Any) -> int:
        cap = self.resource_cap_for_tile(terrain, biome)
        if cap <= 0:
            return 0
        return int(np.clip(np.rint(float(resource_value) * float(cap)), 0, cap))

    def encode_record(self, record: dict[str, Any], flatten: bool = True) -> np.ndarray:
        tiles = record["map"]["tiles"]
        layers = record["map"]["layers"]
        height = self.config.map_height
        width = self.config.map_width
        tensor = np.zeros((height, width, self.feature_dim), dtype=np.float32)

        for y in range(height):
            for x in range(width):
                tile = tiles[y][x]
                terrain = str(tile.get("type", "plain")).lower()
                biome = str(tile.get("biome", "grassland")).lower()
                terrain_idx = self.terrain_to_idx.get(terrain, self.terrain_to_idx["plain"])
                biome_idx = self.biome_to_idx.get(biome, self.biome_to_idx["grassland"])
                tensor[y, x, self.terrain_slice.start + terrain_idx] = 1.0
                tensor[y, x, self.biome_slice.start + biome_idx] = 1.0
                tensor[y, x, self.resource_index] = self.encode_resource_value(
                    terrain,
                    biome,
                    tile.get("resource", 0),
                )
                tensor[y, x, self.base_index] = 1.0 if bool(tile.get("base", False)) else 0.0

        for offset, layer_name in enumerate(self.layer_names):
            layer = np.asarray(layers[layer_name], dtype=np.float32)
            tensor[:, :, self.layer_slice.start + offset] = np.clip(layer, 0.0, 1.0)

        if flatten:
            return tensor.reshape(-1)

        return tensor

    def decode_tensor(
        self,
        tensor: np.ndarray,
        metadata: dict[str, Any] | None = None,
        split: str = "generated",
    ) -> dict[str, Any]:
        grid = self.ensure_grid(tensor)
        terrain_logits = grid[:, :, self.terrain_slice]
        biome_logits = grid[:, :, self.biome_slice]
        resource_values = grid[:, :, self.resource_index]
        base_values = grid[:, :, self.base_index]
        layer_values = grid[:, :, self.layer_slice]
        base_y, base_x = np.unravel_index(np.argmax(base_values), base_values.shape)
        tiles: list[list[dict[str, Any]]] = []
        layers = {
            name: np.clip(layer_values[:, :, idx], 0.0, 1.0).round(4).tolist()
            for idx, name in enumerate(self.layer_names)
        }

        for y in range(self.config.map_height):
            row = []
            for x in range(self.config.map_width):
                terrain_idx = int(np.argmax(terrain_logits[y, x]))
                biome_idx = int(np.argmax(biome_logits[y, x]))
                terrain = self.terrain_types[terrain_idx]
                biome = self._coerce_biome(terrain, self.biome_types[biome_idx])
                is_base = y == base_y and x == base_x
                if is_base:
                    terrain = "base"
                    biome = "colony_core"
                row.append(
                    {
                        "type": terrain,
                        "biome": biome,
                        "resource": self.decode_resource_value(terrain, biome, resource_values[y, x]),
                        "base": bool(is_base),
                    }
                )
            tiles.append(row)

        payload = {
            "metadata": {
                "split": split,
                **(metadata or {}),
            },
            "map": {
                "tiles": tiles,
                "layers": layers,
                "base_position": {"x": int(base_x), "y": int(base_y)},
            },
        }
        return payload

    def reconstruction_metrics(self, original: np.ndarray, reconstructed: np.ndarray) -> dict[str, float]:
        x = self.ensure_grid(original)
        y = self.ensure_grid(reconstructed)
        terrain_acc = float(
            np.mean(
                np.argmax(x[:, :, self.terrain_slice], axis=-1)
                == np.argmax(y[:, :, self.terrain_slice], axis=-1)
            )
        )
        biome_acc = float(
            np.mean(
                np.argmax(x[:, :, self.biome_slice], axis=-1)
                == np.argmax(y[:, :, self.biome_slice], axis=-1)
            )
        )
        base_acc = float(
            np.mean(
                (x[:, :, self.base_index] >= 0.5)
                == (y[:, :, self.base_index] >= 0.5)
            )
        )
        resource_mae = float(
            np.mean(
                np.abs(
                    x[:, :, self.resource_index] - y[:, :, self.resource_index]
                )
            )
        )
        return {
            "terrain_accuracy": terrain_acc,
            "biome_accuracy": biome_acc,
            "base_accuracy": base_acc,
            "resource_mae": resource_mae,
        }

    def ensure_grid(self, tensor: np.ndarray) -> np.ndarray:
        array = np.asarray(tensor, dtype=np.float32)
        if array.ndim == 1:
            return array.reshape(self.config.map_height, self.config.map_width, self.feature_dim)
        if array.ndim == 2 and array.shape[-1] == self.feature_dim:
            return array.reshape(self.config.map_height, self.config.map_width, self.feature_dim)
        if array.ndim == 3:
            return array
        raise ValueError(f"Unsupported tensor shape: {array.shape}")

    def _coerce_biome(self, terrain: str, biome: str) -> str:
        if terrain == "base":
            return "colony_core"
        if terrain == "water":
            return biome if biome in {"lake", "river"} else "lake"
        if terrain == "beach":
            return "coast"
        if terrain == "mountain":
            return "mountain_range"
        if terrain == "sand":
            return "desert"
        if terrain == "forest":
            return biome if biome in {"dense_forest", "young_forest", "wetland"} else "young_forest"
        return biome if biome in {"grassland", "highland", "wetland"} else "grassland"