parabellum 0.4.0__py3-none-any.whl → 0.5.13__py3-none-any.whl

parabellum/__init__.py

@@ -1,22 +1,8 @@
-from .env import Environment, Scenario, make_scenario, State
-from .vis import Visualizer, Skin
-from .gun import bullet_fn
-from . import vis
-from . import terrain_db
-from . import env
-from . import tps
-# from .run import run
+from . import aid, env, geo, types
 
 __all__ = [
+    "geo",
     "env",
-    "terrain_db",
-    "vis",
-    "tps",
-    "Environment",
-    "Scenario",
-    "make_scenario",
-    "State",
-    "Visualizer",
-    "Skin",
-    "bullet_fn",
+    "aid",
+    "types",
 ]

parabellum/aid.py

@@ -3,10 +3,9 @@
 # by: Noah Syrkis
 
 # imports
-import os
 from collections import namedtuple
-from typing import Tuple
 import cartopy.crs as ccrs
+import jax.numpy as jnp
 
 # types
 BBox = namedtuple("BBox", ["north", "south", "east", "west"])  # type: ignore
@@ -23,5 +22,20 @@ def to_mercator(bbox: BBox) -> BBox:
 def to_platecarree(bbox: BBox) -> BBox:
     proj = ccrs.PlateCarree()
     west, south = proj.transform_point(bbox.west, bbox.south, ccrs.Mercator())
+
     east, north = proj.transform_point(bbox.east, bbox.north, ccrs.Mercator())
     return BBox(north=north, south=south, east=east, west=west)
+
+
+def obstacle_mask_fn(limit):
+    def aux(i, j):
+        xs = jnp.linspace(0, i + 1, i + j + 1)
+        ys = jnp.linspace(0, j + 1, i + j + 1)
+        cc = jnp.stack((xs, ys)).astype(jnp.int8)
+        mask = jnp.zeros((limit, limit)).at[*cc].set(1)
+        return mask
+
+    x = jnp.repeat(jnp.arange(limit), limit)
+    y = jnp.tile(jnp.arange(limit), limit)
+    mask = jnp.stack([aux(*c) for c in jnp.stack((x, y)).T])
+    return mask.astype(jnp.int8).reshape(limit, limit, limit, limit)

parabellum/env.py

@@ -1,571 +1,113 @@
-"""Parabellum environment based on SMAX"""
+# env.py
+#   parabellum env
+# by: Noah Syrkis
 
+# % Imports
 import jax.numpy as jnp
-import jax
-from jax import random, Array, vmap, jit
-from flax.struct import dataclass
-import chex
-from jaxmarl.environments.smax.smax_env import SMAX
-
-from math import ceil
-
-from typing import Tuple, Dict, cast
+from jax import random, Array, lax, vmap, debug
+import jax.numpy.linalg as la
+from typing import Tuple
 from functools import partial
-from parabellum import tps, geo, terrain_db
-
-
-@dataclass
-class Scenario:
-    """Parabellum scenario"""
-
-    place: str
-    terrain: tps.Terrain
-    unit_starting_sectors: jnp.ndarray  # must be of size (num_units, 4) where sectors[i] = (x, y, width, height) of the ith unit's spawning sector (in % of the real map)
-    unit_types: chex.Array
-    num_allies: int
-    num_enemies: int
-
-    smacv2_position_generation: bool = False
-    smacv2_unit_type_generation: bool = False
-
-
-@dataclass
-class State:
-    # terrain: Array
-    unit_positions: Array  # fsfds
-    unit_alive: Array
-    unit_teams: Array
-    unit_health: Array
-    unit_types: Array
-    unit_weapon_cooldowns: Array
-    prev_movement_actions: Array
-    prev_attack_actions: Array
-    time: int
-    terminal: bool
-
-
-def make_scenario(
-    place,
-    size,
-    unit_starting_sectors,
-    allies_type,
-    n_allies,
-    enemies_type,
-    n_enemies,
-):
-    if place in terrain_db.db:
-        terrain = terrain_db.make_terrain(terrain_db.db[place], size)
-    else:
-        terrain = geo.geography_fn(place, size)
-    if type(unit_starting_sectors) == list:
-        default_sector = [
-            0,
-            0,
-            size,
-            size,
-        ]  # Noah feel confident that this is right. This means 50% chance. Sorry timothee if you end up here later. my bad bro.
-        correct_unit_starting_sectors = []
-        for i in range(n_allies + n_enemies):
-            selected_sector = None
-            for unit_ids, sector in unit_starting_sectors:
-                if i in unit_ids:
-                    selected_sector = sector
-            if selected_sector is None:
-                selected_sector = default_sector
-            correct_unit_starting_sectors.append(selected_sector)
-        unit_starting_sectors = correct_unit_starting_sectors
-    if type(allies_type) == int:
-        allies = [allies_type] * n_allies
-    else:
-        assert len(allies_type) == n_allies
-        allies = allies_type
-
-    if type(enemies_type) == int:
-        enemies = [enemies_type] * n_enemies
-    else:
-        assert len(enemies_type) == n_enemies
-        enemies = enemies_type
-    unit_types = jnp.array(allies + enemies, dtype=jnp.uint8)
-    return Scenario(
-        place,
-        terrain,
-        unit_starting_sectors,  # type: ignore
-        unit_types,
-        n_allies,
-        n_enemies,
-    )
-
-
-def scenario_fn(place, size):
-    # scenario function for Noah, cos the one above is confusing
-    terrain = geo.geography_fn(place, size)
-    num_allies = 10
-    num_enemies = 10
-    unit_types = jnp.array([0] * num_allies + [1] * num_enemies, dtype=jnp.uint8)
-    # start units in default sectors
-    unit_starting_sectors = jnp.array([[0, 0, 1, 1]] * (num_allies + num_enemies))
-    return Scenario(
-        place=place,
-        terrain=terrain,
-        unit_starting_sectors=unit_starting_sectors,
-        unit_types=unit_types,
-        num_allies=num_allies,
-        num_enemies=num_enemies,
-    )
-
-
-def spawn_fn(rng: jnp.ndarray, units_spawning_sectors):
-    """Spawns n agents on a map."""
-    spawn_positions = []
-    for sector in units_spawning_sectors:
-        rng, key_start, key_noise = random.split(rng, 3)
-        noise = 0.25 + random.uniform(key_noise, (2,)) * 0.5
-        idx = random.choice(key_start, sector[0].shape[0])
-        coord = jnp.array([sector[0][idx], sector[1][idx]])
-        spawn_positions.append(coord + noise)
-    return jnp.array(spawn_positions, dtype=jnp.float32)
-
-
-def sectors_fn(sectors: jnp.ndarray, invalid_spawn_areas: jnp.ndarray):
-    """
-    sectors must be of size (num_units, 4) where sectors[i] = (x, y, width, height) of the ith unit's spawning sector (in % of the real map)
-    """
-    width, height = invalid_spawn_areas.shape
-    spawning_sectors = []
-    for sector in sectors:
-        coordx, coordy = (
-            jnp.array(sector[0] * width, dtype=jnp.int32),
-            jnp.array(sector[1] * height, dtype=jnp.int32),
-        )
-        sector = (
-            invalid_spawn_areas[
-                coordx : coordx + ceil(sector[2] * width),
-                coordy : coordy + ceil(sector[3] * height),
-            ]
-            == 0
-        )
-        valid = jnp.nonzero(sector)
-        if valid[0].shape[0] == 0:
-            raise ValueError(f"Sector {sector} only contains invalid spawn areas.")
-        spawning_sectors.append(
-            jnp.array(valid) + jnp.array([coordx, coordy]).reshape((2, -1))
-        )
-    return spawning_sectors
-
-
-class Environment(SMAX):
-    def __init__(self, scenario: Scenario, **kwargs):
-        map_height, map_width = scenario.terrain.building.shape
-        args = dict(scenario=scenario, map_height=map_height, map_width=map_width)
-        if "unit_type_pushable" in kwargs:
-            self.unit_type_pushable = kwargs["unit_type_pushable"]
-            del kwargs["unit_type_pushable"]
-        else:
-            self.unit_type_pushable = jnp.array([1, 1, 0, 0, 0, 1])
-        if "reset_when_done" in kwargs:
-            self.reset_when_done = kwargs["reset_when_done"]
-            del kwargs["reset_when_done"]
-        else:
-            self.reset_when_done = True
-        super(Environment, self).__init__(**args, walls_cause_death=False, **kwargs)
-        self.terrain = scenario.terrain
-        self.unit_starting_sectors = scenario.unit_starting_sectors
-        # self.unit_type_names = ["tinker", "tailor", "soldier", "spy"]
-        # self.unit_type_health = jnp.array([100, 100, 100, 100], dtype=jnp.float32)
-        # self.unit_type_damage = jnp.array([10, 10, 10, 10], dtype=jnp.float32)
-        self.scenario = scenario
-        self.unit_type_velocities = (
-            jnp.array([3.15, 2.25, 4.13, 3.15, 4.13, 3.15]) / 2.5
-            if "unit_type_velocities" not in kwargs
-            else kwargs["unit_type_velocities"]
-        )
-        self.unit_type_attack_blasts = jnp.zeros((3,), dtype=jnp.float32)  # TODO: add
-        self.max_steps = 200
-        self._push_units_away = lambda state, firmness=1: state  # overwrite push units
-        self.spawning_sectors = sectors_fn(
-            self.unit_starting_sectors,
-            scenario.terrain.building + scenario.terrain.water,
-        )
-        self.resolution = (
-            jnp.array(jnp.max(self.unit_type_sight_ranges), dtype=jnp.int32) * 2
-        )
-        self.t = jnp.tile(jnp.linspace(0, 1, self.resolution), (2, 1))
-
-    def reset(self, rng: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]:  # type: ignore
-        """Environment-specific reset."""
-        unit_positions = spawn_fn(rng, self.spawning_sectors)
-        unit_teams = jnp.zeros((self.num_agents,))
-        unit_teams = unit_teams.at[self.num_allies :].set(1)
-        unit_weapon_cooldowns = jnp.zeros((self.num_agents,))
-        # default behaviour spawn all marines
-        unit_types = cast(Array, self.scenario.unit_types)
-        unit_health = self.unit_type_health[unit_types]
-        state = State(
-            unit_positions=unit_positions,
-            unit_alive=jnp.ones((self.num_agents,), dtype=jnp.bool_),
-            unit_teams=unit_teams,
-            unit_health=unit_health,
-            unit_types=unit_types,
-            prev_movement_actions=jnp.zeros((self.num_agents, 2)),
-            prev_attack_actions=jnp.zeros((self.num_agents,), dtype=jnp.int32),
-            time=0,
-            terminal=False,
-            unit_weapon_cooldowns=unit_weapon_cooldowns,
-            # terrain=self.terrain,
-        )
-        state = self._push_units_away(state)  # type: ignore  could be slow
-        obs = self.get_obs(state)
-        # remove world_state from obs
-        world_state = self.get_world_state(state)
-        obs["world_state"] = jax.lax.stop_gradient(world_state)
-        return obs, state
-
-        # def step_env(self, rng, state: State, action: Array):  # type: ignore
-        # obs, state, rewards, dones, infos = super().step_env(rng, state, action)
-        # delete world_state from obs
-        # <<<<<<< HEAD
-        # obs.pop("world_state")
-        # if not self.reset_when_done:
-        # for key in dones.keys():
-        # dones[key] = False
-        # return obs, state, rewards, dones, infos
-        # =======
-        obs.pop("world_state")
-        if not self.reset_when_done:
-            for key in dones.keys():
-                infos[key] = dones[key]
-                dones[key] = False
-        return obs, state, rewards, dones, infos
-
-    # >>>>>>> aeb13033e57083cc512a60f8f60a3db47a65ac32
-
-    def get_obs_unit_list(self, state: State) -> Dict[str, chex.Array]:  # type: ignore
-        """Applies observation function to state."""
-
-        def get_features(i, j):
-            """Get features of unit j as seen from unit i"""
-            # Can just keep them symmetrical for now.
-            # j here means 'the jth unit that is not i'
-            # The observation is such that allies are always first
-            # so for units in the second team we count in reverse.
-            j = jax.lax.cond(
-                i < self.num_allies, lambda: j, lambda: self.num_agents - j - 1
-            )
-            offset = jax.lax.cond(i < self.num_allies, lambda: 1, lambda: -1)
-            j_idx = jax.lax.cond(
-                ((j < i) & (i < self.num_allies)) | ((j > i) & (i >= self.num_allies)),
-                lambda: j,
-                lambda: j + offset,
-            )
-            empty_features = jnp.zeros(shape=(len(self.unit_features),))
-            features = self._observe_features(state, i, j_idx)
-            visible = (
-                jnp.linalg.norm(state.unit_positions[j_idx] - state.unit_positions[i])
-                < self.unit_type_sight_ranges[state.unit_types[i]]
-            )
-            return jax.lax.cond(
-                visible
-                & state.unit_alive[i]
-                & state.unit_alive[j_idx]
-                & self.has_line_of_sight(
-                    state.unit_positions[j_idx],
-                    state.unit_positions[i],
-                    self.terrain.building + self.terrain.forest,
-                ),
-                lambda: features,
-                lambda: empty_features,
-            )
-
-        get_all_features_for_unit = jax.vmap(get_features, in_axes=(None, 0))
-        get_all_features = jax.vmap(get_all_features_for_unit, in_axes=(0, None))
-        other_unit_obs = get_all_features(
-            jnp.arange(self.num_agents), jnp.arange(self.num_agents - 1)
-        )
-        other_unit_obs = other_unit_obs.reshape((self.num_agents, -1))
-        get_all_self_features = jax.vmap(self._get_own_features, in_axes=(None, 0))
-        own_unit_obs = get_all_self_features(state, jnp.arange(self.num_agents))
-        obs = jnp.concatenate([other_unit_obs, own_unit_obs], axis=-1)
-        return {agent: obs[self.agent_ids[agent]] for agent in self.agents}
-
-    def _our_push_units_away(
-        self, pos, unit_types, firmness: float = 1.0
-    ):  # copy of SMAX._push_units_away but used without state and called inside _world_step to allow more obstacles constraints
-        delta_matrix = pos[:, None] - pos[None, :]
-        dist_matrix = (
-            jnp.linalg.norm(delta_matrix, axis=-1)
-            + jnp.identity(self.num_agents)
-            + 1e-6
-        )
-        radius_matrix = (
-            self.unit_type_radiuses[unit_types][:, None]
-            + self.unit_type_radiuses[unit_types][None, :]
-        )
-        overlap_term = jax.nn.relu(radius_matrix / dist_matrix - 1.0)
-        unit_positions = (
-            pos
-            + firmness * jnp.sum(delta_matrix * overlap_term[:, :, None], axis=1) / 2
-        )
-        return jnp.where(
-            self.unit_type_pushable[unit_types][:, None], unit_positions, pos
-        )
-
-    def has_line_of_sight(self, source, target, raster_input):
-        # suppose that the target is in sight_range of source, otherwise the line of sight might miss some cells
-        cells = jnp.array(
-            source[:, jnp.newaxis] * self.t + (1 - self.t) * target[:, jnp.newaxis],
-            dtype=jnp.int32,
-        )
-        mask = jnp.zeros(raster_input.shape).at[cells[0, :], cells[1, :]].set(1)
-        flag = ~jnp.any(jnp.logical_and(mask, raster_input))
-        return flag
-
-    @partial(jax.jit, static_argnums=(0,))  # replace the _world_step method
-    def _world_step(  # modified version of JaxMARL's SMAX _world_step
-        self,
-        key: chex.PRNGKey,
-        state: State,
-        actions: Tuple[chex.Array, chex.Array],
-    ) -> State:
-        def raster_crossing(pos, new_pos, mask: jnp.ndarray):
-            pos, new_pos = pos.astype(jnp.int32), new_pos.astype(jnp.int32)
-            minimum = jnp.minimum(pos, new_pos)
-            maximum = jnp.maximum(pos, new_pos)
-            mask = jnp.where(jnp.arange(mask.shape[0]) >= minimum[0], mask.T, 0).T
-            mask = jnp.where(jnp.arange(mask.shape[0]) <= maximum[0], mask.T, 0).T
-            mask = jnp.where(jnp.arange(mask.shape[1]) >= minimum[1], mask, 0)
-            mask = jnp.where(jnp.arange(mask.shape[1]) <= maximum[1], mask, 0)
-            return jnp.any(mask)
-
-        def update_position(idx, vec):
-            # Compute the movements slightly strangely.
-            # The velocities below are for diagonal directions
-            # because these are easier to encode as actions than the four
-            # diagonal directions. Then rotate the velocity 45
-            # degrees anticlockwise to compute the movement.
-            pos = cast(Array, state.unit_positions[idx])
-            new_pos = (
-                pos
-                + vec
-                * self.unit_type_velocities[state.unit_types[idx]]
-                * self.time_per_step
-            )
-            # avoid going out of bounds
-            new_pos = jnp.maximum(
-                jnp.minimum(
-                    new_pos, jnp.array([self.map_width - 1, self.map_height - 1])
-                ),
-                jnp.zeros((2,)),
-            )
-
-            #######################################################################
-            ############################################ avoid going into obstacles
-            clash = raster_crossing(
-                pos, new_pos, self.terrain.building + self.terrain.water
-            )
-            new_pos = jnp.where(clash, pos, new_pos)
-
-            #######################################################################
-            #######################################################################
-
-            return new_pos
-
-        #######################################################################
-        ######################################### units close enough to get hit
-
-        def bystander_fn(attacked_idx):
-            idxs = jnp.zeros((self.num_agents,))
-            idxs *= (
-                jnp.linalg.norm(
-                    state.unit_positions - state.unit_positions[attacked_idx], axis=-1
-                )
-                < self.unit_type_attack_blasts[state.unit_types[attacked_idx]]
-            )
-            return idxs
-
-            #######################################################################
-            #######################################################################
-
-        def update_agent_health(idx, action, key):  # TODO: add attack blasts
-            # for team 1, their attack actions are labelled in
-            # reverse order because that is the order they are
-            # observed in
-            attacked_idx = jax.lax.cond(
-                idx < self.num_allies,
-                lambda: action + self.num_allies - self.num_movement_actions,
-                lambda: self.num_allies - 1 - (action - self.num_movement_actions),
-            )
-            attacked_idx = cast(int, attacked_idx)  # Cast to int
-            # deal with no-op attack actions (i.e. agents that are moving instead)
-            attacked_idx = jax.lax.select(
-                action < self.num_movement_actions, idx, attacked_idx
-            )
-            distance = jnp.linalg.norm(
-                state.unit_positions[idx] - state.unit_positions[attacked_idx]
-            )
-            attack_valid = (
-                (distance <= self.unit_type_attack_ranges[state.unit_types[idx]])
-                & state.unit_alive[idx]
-                & state.unit_alive[attacked_idx]
-            )
-            attack_valid = attack_valid & (idx != attacked_idx)
-            attack_valid = attack_valid & (state.unit_weapon_cooldowns[idx] <= 0.0)
-            health_diff = jax.lax.select(
-                attack_valid,
-                -self.unit_type_attacks[state.unit_types[idx]],
-                0.0,
-            )
-            health_diff = jnp.where(
-                state.unit_types[idx] == 1,
-                health_diff
-                * distance
-                / self.unit_type_attack_ranges[state.unit_types[idx]],
-                health_diff,
-            )
-            # design choice based on the pysc2 randomness details.
-            # See https://github.com/deepmind/pysc2/blob/master/docs/environment.md#determinism-and-randomness
-
-            #########################################################
-            ############################### Add bystander health diff
-
-            # bystander_idxs = bystander_fn(attacked_idx)  # TODO: use
-            # bystander_valid = (
-            #     jnp.where(attack_valid, bystander_idxs, jnp.zeros((self.num_agents,)))
-            #     .astype(jnp.bool_)  # type: ignore
-            #     .astype(jnp.float32)
-            # )
-            # bystander_health_diff = (
-            #     bystander_valid * -self.unit_type_attacks[state.unit_types[idx]]
-            # )
-
-            #########################################################
-            #########################################################
-
-            cooldown_deviation = jax.random.uniform(
-                key, minval=-self.time_per_step, maxval=2 * self.time_per_step
-            )
-            cooldown = (
-                self.unit_type_weapon_cooldowns[state.unit_types[idx]]
-                + cooldown_deviation
-            )
-            cooldown_diff = jax.lax.select(
-                attack_valid,
-                # subtract the current cooldown because we are
-                # going to add it back. This way we effectively
-                # set the new cooldown to `cooldown`
-                cooldown - state.unit_weapon_cooldowns[idx],
-                -self.time_per_step,
-            )
-            return (
-                health_diff,
-                attacked_idx,
-                cooldown_diff,
-                # (bystander_health_diff, bystander_idxs),
-            )
-
-        def perform_agent_action(idx, action, key):
-            movement_action, attack_action = action
-            new_pos = update_position(idx, movement_action)
-            health_diff, attacked_idxes, cooldown_diff = update_agent_health(
-                idx, attack_action, key
-            )
-
-            return new_pos, (health_diff, attacked_idxes), cooldown_diff
-
-        keys = jax.random.split(key, num=self.num_agents)
-        pos, (health_diff, attacked_idxes), cooldown_diff = jax.vmap(
-            perform_agent_action
-        )(jnp.arange(self.num_agents), actions, keys)
-
-        # units push each other
-        new_pos = self._our_push_units_away(pos, state.unit_types)
-        clash = jax.vmap(raster_crossing, in_axes=(0, 0, None))(
-            pos, new_pos, self.terrain.building + self.terrain.water
-        )
-        pos = jax.vmap(jnp.where)(clash, pos, new_pos)
-        # avoid going out of bounds
-        pos = jnp.maximum(
-            jnp.minimum(pos, jnp.array([self.map_width - 1, self.map_height - 1])),  # type: ignore
-            jnp.zeros((2,)),
-        )
-
-        # Multiple enemies can attack the same unit.
-        # We have `(health_diff, attacked_idx)` pairs.
-        # `jax.lax.scatter_add` aggregates these exactly
-        # in the way we want -- duplicate idxes will have their
-        # health differences added together. However, it is a
-        # super thin wrapper around the XLA scatter operation,
-        # which has this bonkers syntax and requires this dnums
-        # parameter. The usage here was inferred from a test:
-        # https://github.com/google/jax/blob/main/tests/lax_test.py#L2296
-        dnums = jax.lax.ScatterDimensionNumbers(
-            update_window_dims=(),
-            inserted_window_dims=(0,),
-            scatter_dims_to_operand_dims=(0,),
-        )
-        unit_health = jnp.maximum(
-            jax.lax.scatter_add(
-                state.unit_health,
-                jnp.expand_dims(attacked_idxes, 1),
-                health_diff,
-                dnums,
-            ),
-            0.0,
-        )
-
-        #########################################################
-        ############################ subtracting bystander health
-
-        # _, bystander_health_diff = bystander
-        # unit_health -= bystander_health_diff.sum(axis=0)  # might be axis=1
-
-        #########################################################
-        #########################################################
-
-        unit_weapon_cooldowns = state.unit_weapon_cooldowns + cooldown_diff
-        # replace unit health, unit positions and unit weapon cooldowns
-        state = state.replace(  # type: ignore
-            unit_health=unit_health,
-            unit_positions=pos,
-            unit_weapon_cooldowns=unit_weapon_cooldowns,
-        )
-        return state
-
-
-if __name__ == "__main__":
-    n_envs = 4
-
-    n_allies = 10
-    scenario_kwargs = {
-        "allies_type": 0,
-        "n_allies": n_allies,
-        "enemies_type": 0,
-        "n_enemies": n_allies,
-        "place": "Vesterbro, Copenhagen, Denmark",
-        "size": 100,
-        "unit_starting_sectors": [
-            ([i for i in range(n_allies)], [0.0, 0.45, 0.1, 0.1]),
-            ([n_allies + i for i in range(n_allies)], [0.8, 0.5, 0.1, 0.1]),
-        ],
-    }
-    scenario = make_scenario(**scenario_kwargs)
-    env = Environment(scenario)
-    rng, reset_rng = random.split(random.PRNGKey(0))
-    reset_key = random.split(reset_rng, n_envs)
-    obs, state = vmap(env.reset)(reset_key)
-    state_seq = []
-
-    import time
 
-    step = vmap(jit(env.step))
-    tic = time.time()
-    for i in range(10):
-        rng, act_rng, step_rng = random.split(rng, 3)
-        act_key = random.split(act_rng, (len(env.agents), n_envs))
-        act = {
-            a: vmap(env.action_space(a).sample)(act_key[i])
-            for i, a in enumerate(env.agents)
-        }
-        step_key = random.split(step_rng, n_envs)
-        state_seq.append((step_key, state, act))
-        obs, state, reward, done, infos = step(step_key, state, act)
-        tic = time.time()
+from parabellum.geo import geography_fn
+from parabellum.types import Action, State, Obs, Scene
+from parabellum import aid
+import equinox as eqx
+
+
+# %% Dataclass ################################################################
+class Env:
+    def __init__(self, cfg):
+        self.cfg = cfg
+
+    def reset(self, rng: Array, scene: Scene) -> Tuple[Obs, State]:
+        return init_fn(rng, self, scene)
+
+    def step(self, rng: Array, scene: Scene, state: State, action: Action) -> Tuple[Obs, State]:
+        return obs_fn(self, scene, state), step_fn(rng, self, scene, state, action)
+
+    @property
+    def num_units(self):
+        return sum(self.cfg.counts.allies.values()) + sum(self.cfg.counts.enemies.values())
+
+    @property
+    def num_allies(self):
+        return sum(self.cfg.counts.allies.values())
+
+    @property
+    def num_enemies(self):
+        return sum(self.cfg.counts.enemies.values())
+
+
+# %% Functions ################################################################
+@eqx.filter_jit
+def init_fn(rng: Array, env: Env, scene: Scene) -> Tuple[Obs, State]:  # initialize -----
+    keys = random.split(rng)
+    health = jnp.ones(env.num_units) * scene.unit_type_health[scene.unit_types]
+    pos = random.normal(keys[1], (scene.unit_types.size, 2)) * 2 + env.cfg.size / 2
+    state = State(unit_position=pos, unit_health=health, unit_cooldown=jnp.zeros(env.num_units))  # state --
+    return obs_fn(env, scene, state), state  # return observation and state of agents --
+
+
+@eqx.filter_jit  # knn from env.cfg never changes, so we can jit it
+def obs_fn(env, scene: Scene, state: State) -> Obs:  # return info about neighbors ---
+    distances = la.norm(state.unit_position[:, None] - state.unit_position, axis=-1)  # all dist --
+    dists, idxs = lax.approx_min_k(distances, k=env.cfg.knn)
+    mask = mask_fn(scene, state, dists, idxs)
+    health = state.unit_health[idxs] * mask
+    cooldown = state.unit_cooldown[idxs] * mask
+    unit_pos = (state.unit_position[:, None, ...] - state.unit_position[idxs]) * mask[..., None]
+    return Obs(unit_id=idxs, unit_pos=unit_pos, unit_health=health, unit_cooldown=cooldown)
+
+
+@eqx.filter_jit
+def step_fn(rng, env: Env, scene: Scene, state: State, action: Action) -> State:  # update agents ---
+    newpos = state.unit_position + action.coord * (1 - action.kinds[..., None])
+    bounds = ((newpos < 0).any(axis=-1) | (newpos >= env.cfg.size).any(axis=-1))[..., None]
+    builds = (scene.terrain.building[*newpos.astype(jnp.int32).T] > 0)[..., None]
+    newpos = jnp.where(bounds | builds, state.unit_position, newpos)  # use old pos if new is not valid
+    new_hp = blast_fn(rng, env, scene, state, action)
+    return State(unit_position=newpos, unit_health=new_hp, unit_cooldown=state.unit_cooldown)  # return -
+
+
+def blast_fn(rng, env: Env, scene: Scene, state: State, action: Action):  # update agents ---
+    dist = la.norm(state.unit_position[None, ...] - (state.unit_position + action.coord)[:, None, ...], axis=-1)
+    hits = dist <= scene.unit_type_reach[scene.unit_types][None, ...] * action.kinds[..., None]  # mask non attack act
+    damage = (hits * scene.unit_type_damage[scene.unit_types][None, ...]).sum(axis=-1)
+    return state.unit_health - damage
+
+
+# @eqx.filter_jit
+def scene_fn(cfg):  # init's a scene
+    aux = lambda key: jnp.array([x[key] for x in sorted(cfg.types, key=lambda x: x.name)])  # noqa
+    attrs = ["health", "damage", "reload", "reach", "sight", "speed"]
+    kwargs = {f"unit_type_{a}": aux(a) for a in attrs} | {"terrain": geography_fn(cfg.place, cfg.size)}
+    num_allies, num_enemies = sum(cfg.counts.allies.values()), sum(cfg.counts.enemies.values())
+    unit_teams = jnp.concat((jnp.zeros(num_allies), jnp.ones(num_enemies))).astype(jnp.int32)
+    aux = lambda t: jnp.concat([jnp.zeros(x) + i for i, x in enumerate([x[1] for x in sorted(cfg.counts[t].items())])])  # noqa
+    unit_types = jnp.concat((aux("allies"), aux("enemies"))).astype(jnp.int32)
+    mask = aid.obstacle_mask_fn(max([x["sight"] for x in cfg.types]))
+    return Scene(unit_teams=unit_teams, unit_types=unit_types, mask=mask, **kwargs)  # type: ignore
+
+
+@eqx.filter_jit
+def mask_fn(scene, state, dists, idxs):
+    mask = dists < scene.unit_type_sight[scene.unit_types][..., None]  # mask for removing hidden
+    mask = mask | obstacle_fn(scene, state.unit_position[idxs].astype(jnp.int8))
+    return mask
+
+
+@partial(vmap, in_axes=(None, 0))  # 5 x 2 # not the best name for a fn
+def obstacle_fn(scene, pos):
+    slice = slice_fn(scene, pos[0], pos)
+    return slice
+
+
+@partial(vmap, in_axes=(None, None, 0))
+def slice_fn(scene, source, target):  # returns a 10 x 10 view with unit at top left corner, and terrain downwards
+    delta = ((source - target) >= 0) * 2 - 1
+    coord = jnp.sort(jnp.stack((source, source + delta * 10)), axis=0)[0]
+    slice = lax.dynamic_slice(scene.terrain.building, coord, (scene.mask.shape[-1], scene.mask.shape[-1]))
+    slice = lax.cond(delta[0] == 1, lambda: jnp.flip(slice), lambda: slice)
+    slice = lax.cond(delta[1] == 1, lambda: jnp.flip(slice, axis=1), lambda: slice)
+    return (scene.mask[*jnp.abs(source - target)] * slice).sum() == 0

parabellum/geo.py

@@ -3,25 +3,23 @@
 # by: Noah Syrkis
 
 # %% Imports
-from parabellum import tps
-import rasterio
 from rasterio import features, transform
+from jax import tree
 from geopy.geocoders import Nominatim
 from geopy.distance import distance
 import contextily as cx
 import jax.numpy as jnp
 import cartopy.crs as ccrs
 from jaxtyping import Array
-import numpy as np
 from shapely import box
 import osmnx as ox
 import geopandas as gpd
 from collections import namedtuple
 from typing import Tuple
 import matplotlib.pyplot as plt
-import seaborn as sns
-import os
+from cachier import cachier
 from jax.scipy.signal import convolve
+from parabellum.types import Terrain
 
 # %% Types
 Coords = Tuple[float, float]
@@ -79,27 +77,23 @@ def basemap_fn(bbox: BBox, gdf) -> Array:
     return image
 
 
-def geography_fn(place, buffer=400):
+@cachier()
+def geography_fn(place, buffer) -> Terrain:
     bbox = get_bbox(place, buffer)
     map_data = ox.features_from_bbox(bbox=bbox, tags=tags)
     gdf = gpd.GeoDataFrame(map_data)
-    gdf = gdf.clip(box(bbox.west, bbox.south, bbox.east, bbox.north)).to_crs(
-        "EPSG:3857"
-    )
+    gdf = gdf.clip(box(bbox.west, bbox.south, bbox.east, bbox.north)).to_crs("EPSG:3857")
     raster = raster_fn(gdf, shape=(buffer, buffer))
     basemap = jnp.rot90(basemap_fn(bbox, gdf), 3)
-    # 0: building", 1: "water", 2: "highway", 3: "forest", 4: "garden"
     kernel = jnp.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
-    trans = lambda x: jnp.rot90(x, 3)
-    # <<<<<<< HEAD
-    terrain = tps.Terrain(
+    trans = lambda x: jnp.rot90(x, 3)  # noqa
+    terrain = Terrain(
         building=trans(raster[0]),
-        water=trans(
-            raster[1] - convolve(raster[1] * raster[2], kernel, mode="same") > 0
-        ),
+        water=trans(raster[1] - convolve(raster[1] * raster[2], kernel, mode="same") > 0),
         forest=trans(jnp.logical_or(raster[3], raster[4])),
         basemap=basemap,
     )
+    terrain = tree.map(lambda x: x.astype(jnp.int16), terrain)
     return terrain
 
 
@@ -128,25 +122,25 @@ def feature_fn(t, feature, gdf, shape):
 
 
 # %%
-def normalize(x):
-    return (np.array(x) - m) / (M - m)
-
-
-def get_bridges(gdf):
-    xmin, ymin, xmax, ymax = gdf.total_bounds
-    m = np.array([xmin, ymin])
-    M = np.array([xmax, ymax])
-
-    bridges = {}
-    for idx, bridge in gdf[gdf["bridge"] == "yes"].iterrows():
-        if type(bridge["name"]) == str:
-            bridges[idx[1]] = {
-                "name": bridge["name"],
-                "coords": normalize(
-                    [bridge.geometry.centroid.x, bridge.geometry.centroid.y]
-                ),
-            }
-    return bridges
+# def normalize(x):
+#     return (np.array(x) - m) / (M - m)
+
+
+# def get_bridges(gdf):
+#     xmin, ymin, xmax, ymax = gdf.total_bounds
+#     m = np.array([xmin, ymin])
+#     M = np.array([xmax, ymax])
+
+#     bridges = {}
+#     for idx, bridge in gdf[gdf["bridge"] == "yes"].iterrows():
+#         if type(bridge["name"]) == str:
+#             bridges[idx[1]] = {
+#                 "name": bridge["name"],
+#                 "coords": normalize(
+#                     [bridge.geometry.centroid.x, bridge.geometry.centroid.y]
+#                 ),
+#             }
+#     return bridges
 
 
 """

parabellum/model.py

@@ -0,0 +1,6 @@
+# model.py
+#   jax model for mapping from observation to action
+# by: Noah Syrkis
+
+# Imports
+import jax.numpy as jnp

parabellum/ppo.py

@@ -0,0 +1 @@
+

parabellum/terrain_db.py

@@ -1,22 +1,22 @@
 # %%
 import numpy as np
 import jax.numpy as jnp
-from parabellum import tps
+from parabellum.types import Terrain
 
 
 # %%
 def map_raster_from_line(raster, line, size):
     x0, y0, dx, dy = line
-    x0 = int(x0*size)
-    y0 = int(y0*size)
-    dx = int(dx*size)
-    dy = int(dy*size)
+    x0 = int(x0 * size)
+    y0 = int(y0 * size)
+    dx = int(dx * size)
+    dy = int(dy * size)
     max_T = int(2**0.5 * size)
-    for t in range(max_T+1):
-        alpha = t/float(max_T)
-        x = x0 if dx == 0 else int((1 - alpha) * x0 + alpha * (x0+dx))
-        y = y0 if dy == 0 else  int((1 - alpha) * y0 + alpha * (y0+dy))
-        if 0<=x<size and 0<=y<size:
+    for t in range(max_T + 1):
+        alpha = t / float(max_T)
+        x = x0 if dx == 0 else int((1 - alpha) * x0 + alpha * (x0 + dx))
+        y = y0 if dy == 0 else int((1 - alpha) * y0 + alpha * (y0 + dy))
+        if 0 <= x < size and 0 <= y < size:
             raster[x, y] = 1
     return raster
 
@@ -24,20 +24,21 @@ def map_raster_from_line(raster, line, size):
 # %%
 def map_raster_from_rect(raster, rect, size):
     x0, y0, dx, dy = rect
-    x0 = int(x0*size)
-    y0 = int(y0*size)
-    dx = int(dx*size)
-    dy = int(dy*size)
-    raster[x0:x0+dx, y0:y0+dy] = 1
+    x0 = int(x0 * size)
+    y0 = int(y0 * size)
+    dx = int(dx * size)
+    dy = int(dy * size)
+    raster[x0 : x0 + dx, y0 : y0 + dy] = 1
     return raster
 
 
 # %%
-building_color = jnp.array([201,199,198, 255])
+building_color = jnp.array([201, 199, 198, 255])
 water_color = jnp.array([193, 237, 254, 255])
-forest_color = jnp.array([197,214,185, 255])
+forest_color = jnp.array([197, 214, 185, 255])
 empty_color = jnp.array([255, 255, 255, 255])
 
+
 def make_terrain(terrain_args, size):
     args = {}
     for key, config in terrain_args.items():
@@ -49,44 +50,75 @@ def make_terrain(terrain_args, size):
                 elif "rect" in elem:
                     raster = map_raster_from_rect(raster, elem["rect"], size)
         args[key] = jnp.array(raster.T)
-    basemap = jnp.where(args["building"][:,:,None], jnp.tile(building_color, (size, size, 1)), jnp.tile(empty_color, (size,size, 1)))
-    basemap = jnp.where(args["water"][:,:,None], jnp.tile(water_color, (size, size, 1)), basemap)
-    basemap = jnp.where(args["forest"][:,:,None], jnp.tile(forest_color, (size, size, 1)), basemap)
+    basemap = jnp.where(
+        args["building"][:, :, None], jnp.tile(building_color, (size, size, 1)), jnp.tile(empty_color, (size, size, 1))
+    )
+    basemap = jnp.where(args["water"][:, :, None], jnp.tile(water_color, (size, size, 1)), basemap)
+    basemap = jnp.where(args["forest"][:, :, None], jnp.tile(forest_color, (size, size, 1)), basemap)
     args["basemap"] = basemap
-    return tps.Terrain(**args)
+    return Terrain(**args)
 
 
 # %%
 db = {
-    "blank": {'building': None, 'water': None, 'forest': None},
-    "F": {'building': [{"line": [0.25, 0.33, 0.5, 0]}, {"line":[0.75, 0.33, 0., 0.25]}, {"line":[0.50, 0.33, 0., 0.25]}], 'water': None, 'forest': None},
-    "stronghold": {'building': [
-    {"line":[0.2, 0.275, 0.2, 0.]}, {"line":[0.2, 0.275, 0.0, 0.2]}, 
-    {"line":[0.4, 0.275, 0.0, 0.2]}, {"line":[0.2, 0.475, 0.2, 0.]}, 
-
-    {"line":[0.2, 0.525, 0.2, 0.]}, {"line": [0.2, 0.525, 0.0, 0.2]}, 
-    {"line":[0.4, 0.525, 0.0, 0.2]}, {"line": [0.2, 0.725, 0.525, 0.]}, 
-
-    {"line":[0.75, 0.25, 0., 0.2]}, {"line":[0.75, 0.55, 0., 0.19]},
-    {"line":[0.6, 0.25, 0.15, 0.]}], 'water': None, 'forest': None},
-    "playground": {'building': [{"line":[0.5, 0.5, 0.5, 0.]}], 'water': None, 'forest': None},
+    "blank": {"building": None, "water": None, "forest": None},
+    "F": {
+        "building": [
+            {"line": [0.25, 0.33, 0.5, 0]},
+            {"line": [0.75, 0.33, 0.0, 0.25]},
+            {"line": [0.50, 0.33, 0.0, 0.25]},
+        ],
+        "water": None,
+        "forest": None,
+    },
+    "stronghold": {
+        "building": [
+            {"line": [0.2, 0.275, 0.2, 0.0]},
+            {"line": [0.2, 0.275, 0.0, 0.2]},
+            {"line": [0.4, 0.275, 0.0, 0.2]},
+            {"line": [0.2, 0.475, 0.2, 0.0]},
+            {"line": [0.2, 0.525, 0.2, 0.0]},
+            {"line": [0.2, 0.525, 0.0, 0.2]},
+            {"line": [0.4, 0.525, 0.0, 0.2]},
+            {"line": [0.2, 0.725, 0.525, 0.0]},
+            {"line": [0.75, 0.25, 0.0, 0.2]},
+            {"line": [0.75, 0.55, 0.0, 0.19]},
+            {"line": [0.6, 0.25, 0.15, 0.0]},
+        ],
+        "water": None,
+        "forest": None,
+    },
+    "playground": {"building": [{"line": [0.5, 0.5, 0.5, 0.0]}], "water": None, "forest": None},
     "playground2": {
-        'building': [], 
-        "water": [{"rect":[0., 0.8, 0.1, 0.1]}, {"rect": [0.2, 0.8, 0.8, 0.1]}], 
-        "forest": [{"rect": [0., 0., 1., 0.2]}]
+        "building": [],
+        "water": [{"rect": [0.0, 0.8, 0.1, 0.1]}, {"rect": [0.2, 0.8, 0.8, 0.1]}],
+        "forest": [{"rect": [0.0, 0.0, 1.0, 0.2]}],
+    },
+    "triangle": {
+        "building": [{"line": [0.33, 0.0, 0.0, 1.0]}, {"line": [0.66, 0.0, 0.0, 1.0]}],
+        "water": None,
+        "forest": None,
     },
-    "triangle":  {'building': [{"line": [0.33, 0., 0., 1.]}, {"line": [0.66, 0., 0., 1.]}], 'water': None, 'forest': None},
     "u_shape": {
-        'building': [], 
-        "water": [{"rect": [0.15, 0.2, 0.1, 0.5]}, {"rect": [0.4, 0.2, 0.1, 0.5]}, {"rect": [0.2, 0.2, 0.25, 0.1]}], 
-        "forest": []
+        "building": [],
+        "water": [{"rect": [0.15, 0.2, 0.1, 0.5]}, {"rect": [0.4, 0.2, 0.1, 0.5]}, {"rect": [0.2, 0.2, 0.25, 0.1]}],
+        "forest": [],
     },
     "bridges": {
-        'building': [], 
-        "water": [{"rect": [0.475, 0., 0.05, 0.1]}, {"rect": [0.475, 0.15, 0.05, 0.575]}, {"rect": [0.475, 0.775, 0.05, 1.]},
-                  {"rect": [0., 0.475, 0.225, 0.05]}, {"rect": [0.275, 0.475, 0.45, 0.05]}, {"rect": [0.775, 0.475, 0.23, 0.05]}], 
-        "forest": [{"rect": [0.1, 0.625, 0.275, 0.275]}, {"rect": [0.725, 0., 0.3, 0.275]}, ]
-    }
+        "building": [],
+        "water": [
+            {"rect": [0.475, 0.0, 0.05, 0.1]},
+            {"rect": [0.475, 0.15, 0.05, 0.575]},
+            {"rect": [0.475, 0.775, 0.05, 1.0]},
+            {"rect": [0.0, 0.475, 0.225, 0.05]},
+            {"rect": [0.275, 0.475, 0.45, 0.05]},
+            {"rect": [0.775, 0.475, 0.23, 0.05]},
+        ],
+        "forest": [
+            {"rect": [0.1, 0.625, 0.275, 0.275]},
+            {"rect": [0.725, 0.0, 0.3, 0.275]},
+        ],
+    },
 }
 
 # %% [raw]
@@ -128,7 +160,7 @@ if __name__ == "__main__":
     plt.imshow(jnp.rot90(terrain.basemap))
     bl = (39.5, 5)
     tr = (44.5, 10)
-    plt.scatter(bl[0], 49-bl[1])
-    plt.scatter(tr[0], 49-tr[1], marker="+")
+    plt.scatter(bl[0], 49 - bl[1])
+    plt.scatter(tr[0], 49 - tr[1], marker="+")
 
 # %%

parabellum/types.py

@@ -0,0 +1,54 @@
+# types.py
+#   parabellum types
+# by: Noah Syrkis
+
+# imports
+from chex import dataclass
+from jaxtyping import Array, Bool, Float16
+
+
+# dataclasses
+@dataclass
+class State:
+    unit_position: Array
+    unit_health: Array
+    unit_cooldown: Array
+
+
+@dataclass
+class Obs:
+    unit_id: Array
+    unit_pos: Array
+    unit_health: Array
+    unit_cooldown: Array
+
+
+@dataclass
+class Action:
+    coord: Float16[Array, "... 2"]  # noqa
+    kinds: Bool[Array, "..."]
+
+
+@dataclass
+class Terrain:
+    building: Array
+    water: Array
+    forest: Array
+    basemap: Array
+
+
+@dataclass
+class Scene:
+    terrain: Terrain
+    mask: Array
+
+    unit_types: Array
+    unit_teams: Array
+
+    unit_type_health: Array
+    unit_type_damage: Array
+    unit_type_reload: Array
+
+    unit_type_reach: Array
+    unit_type_sight: Array
+    unit_type_speed: Array

{parabellum-0.4.0.dist-info → parabellum-0.5.13.dist-info}/METADATA

@@ -1,39 +1,41 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: parabellum
-Version: 0.4.0
+Version: 0.5.13
 Summary: Parabellum environment for parallel warfare simulation
-Home-page: https://github.com/syrkis/parabellum
-License: MIT
-Keywords: warfare,simulation,parallel,environment
 Author: Noah Syrkis
 Author-email: desk@syrkis.com
 Requires-Python: >=3.11,<3.12
-Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: brax (>=0.12.1,<0.13.0)
+Requires-Dist: cachier (>=3.1.2,<4.0.0)
 Requires-Dist: cartopy (>=0.23.0,<0.24.0)
 Requires-Dist: contextily (>=1.6.0,<2.0.0)
-Requires-Dist: darkdetect (>=0.8.0,<0.9.0)
+Requires-Dist: distrax (>=0.1.5,<0.2.0)
 Requires-Dist: einops (>=0.8.0,<0.9.0)
+Requires-Dist: equinox (>=0.11.11,<0.12.0)
+Requires-Dist: evosax (>=0.1.6,<0.2.0)
+Requires-Dist: flashbax (>=0.1.2,<0.2.0)
+Requires-Dist: flax (>=0.10.4,<0.11.0)
 Requires-Dist: folium (>=0.17.0,<0.18.0)
 Requires-Dist: geopy (>=2.4.1,<3.0.0)
+Requires-Dist: gymnax (>=0.0.8,<0.0.9)
 Requires-Dist: ipykernel (>=6.29.5,<7.0.0)
-Requires-Dist: jax (==0.4.17)
-Requires-Dist: jaxmarl (==0.0.3)
+Requires-Dist: jax (>=0.5.0,<0.6.0)
+Requires-Dist: jax-tqdm (>=0.3.1,<0.4.0)
 Requires-Dist: jaxtyping (>=0.2.33,<0.3.0)
 Requires-Dist: jupyterlab (>=4.2.2,<5.0.0)
-Requires-Dist: moviepy (>=1.0.3,<2.0.0)
-Requires-Dist: numpy (<2)
-Requires-Dist: opencv-python (>=4.10.0.84,<5.0.0.0)
+Requires-Dist: navix (>=0.7.0,<0.8.0)
+Requires-Dist: numpy (>=2.2.3,<3.0.0)
+Requires-Dist: omegaconf (>=2.3.0,<3.0.0)
+Requires-Dist: optax (>=0.2.4,<0.3.0)
 Requires-Dist: osmnx (==2.0.0b0)
 Requires-Dist: pandas (>=2.2.2,<3.0.0)
 Requires-Dist: poetry (>=1.8.3,<2.0.0)
-Requires-Dist: pygame (>=2.5.2,<3.0.0)
 Requires-Dist: rasterio (>=1.3.10,<2.0.0)
-Requires-Dist: seaborn (>=0.13.2,<0.14.0)
 Requires-Dist: stadiamaps (>=3.2.1,<4.0.0)
 Requires-Dist: tqdm (>=4.66.4,<5.0.0)
-Project-URL: Repository, https://github.com/syrkis/parabellum
+Requires-Dist: wandb (>=0.19.7,<0.20.0)
 Description-Content-Type: text/markdown
 
 # Parabellum

parabellum-0.5.13.dist-info/RECORD

@@ -0,0 +1,15 @@
+parabellum/__init__.py,sha256=Og0bpKlQtkWCJ1yaQk98LniI1sG3B7GAR7aPMFC-v74,96
+parabellum/aid.py,sha256=hCp-eDONcloKMNpYni61cE6St_67Lks2ivS5eJjSHFQ,1379
+parabellum/env.py,sha256=NmDXpjzwun1PsgaHPcIH-i4H8TfkY0Iv1sEy2PKoMnA,5357
+parabellum/geo.py,sha256=Zv2GCwaYKsfB6CE_cwEHhddU9p9Z5e5lL16iNfETVmw,6181
+parabellum/gun.py,sha256=nvsJdcZ2Qd6lbPlAgsUiaLhstTi1UdLQ8kOnbCenucY,2618
+parabellum/model.py,sha256=o40jW2vp3Fwxt1KqykZ-qZXs73-H24nA4iiRoetAobA,117
+parabellum/pcg.py,sha256=d8KC_lbc4WUUUPaTdPJSx27VMGioys3jSGOWJ-2EahU,968
+parabellum/ppo.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+parabellum/run.py,sha256=Q53__AxzROZNgfZLVU5LDdcT61UMCkmQ_Q5wWUIrnqo,3473
+parabellum/terrain_db.py,sha256=lPd56Qe4_xuVFJcKsftE9RvNN9nPr69GzBJZIS5nKrY,5207
+parabellum/types.py,sha256=WwMzSQ5qnRo6rqKTHIGlWPcg57lx6ENDSK5DFr3R7-s,832
+parabellum/vis.py,sha256=ABHveJj0fLRWkxOv3LFIXK20QtdGhjskuFLsp7iTFu0,6185
+parabellum-0.5.13.dist-info/METADATA,sha256=g_tMs-pztJXVRE__E9VLORyZANegyalSoT7ilMz5oRc,2772
+parabellum-0.5.13.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+parabellum-0.5.13.dist-info/RECORD,,

{parabellum-0.4.0.dist-info → parabellum-0.5.13.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.9.0
+Generator: poetry-core 2.1.2
 Root-Is-Purelib: true
 Tag: py3-none-any

parabellum/tps.py

@@ -1,17 +0,0 @@
-# %%
-# tps.py
-#   parabellum types and dataclasses
-# by: Noah Syrkis
-
-# %% Imports
-from chex import dataclass
-from jaxtyping import Array
-
-
-# %% Dataclasses
-@dataclass
-class Terrain:
-    building: Array
-    water: Array
-    forest: Array
-    basemap: Array

parabellum-0.4.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-parabellum/__init__.py,sha256=hIOLir7wgaf_HU4j8uos7PaCrofqPQcr3FcMlBsZyr8,406
-parabellum/aid.py,sha256=BPabjN4BUq1HRhkwbc9pCNsXSF_ALiG8W8cHWTWeEH4,900
-parabellum/env.py,sha256=2FDOI90IuoOTFV5DhLLWWMpuaj4mcwoPup24KK-duYI,22907
-parabellum/geo.py,sha256=PJs9UevibuokDVb3oJWNHvYHlMYGCxB5OkNSbDj48vI,6198
-parabellum/gun.py,sha256=nvsJdcZ2Qd6lbPlAgsUiaLhstTi1UdLQ8kOnbCenucY,2618
-parabellum/pcg.py,sha256=d8KC_lbc4WUUUPaTdPJSx27VMGioys3jSGOWJ-2EahU,968
-parabellum/run.py,sha256=Q53__AxzROZNgfZLVU5LDdcT61UMCkmQ_Q5wWUIrnqo,3473
-parabellum/terrain_db.py,sha256=5lHzbX94lzkb--cETpraXS42G4T4tKekISpTm4yaYEE,4748
-parabellum/tps.py,sha256=of-RBdelAbNCHQZd1I22RWmZkwUEh6f161mx0X_G2tE,257
-parabellum/vis.py,sha256=ABHveJj0fLRWkxOv3LFIXK20QtdGhjskuFLsp7iTFu0,6185
-parabellum-0.4.0.dist-info/METADATA,sha256=GrTOPfKE1HHK2R3RR7gpYgLBFzSIgOgDTpt7h3zFlCM,2707
-parabellum-0.4.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-parabellum-0.4.0.dist-info/RECORD,,