parabellum 0.2.25__py3-none-any.whl → 0.3.0__py3-none-any.whl

parabellum/__init__.py

@@ -1,18 +1,22 @@
-from .env import Environment, Scenario, scenarios, make_scenario, State
+from .env import Environment, Scenario, make_scenario, State
 from .vis import Visualizer, Skin
-from .map import terrain_fn
 from .gun import bullet_fn
-# from .aid import aid
+from . import vis
+from . import map
+from . import env
+from . import tps
 # from .run import run
 
 __all__ = [
+    "env",
+    "map",
+    "vis",
+    "tps",
     "Environment",
     "Scenario",
-    "scenarios",
     "make_scenario",
     "State",
     "Visualizer",
     "Skin",
-    "terrain_fn",
     "bullet_fn",
 ]

parabellum/aid.py

@@ -3,3 +3,25 @@
 # by: Noah Syrkis
 
 # imports
+import os
+from collections import namedtuple
+from typing import Tuple
+import cartopy.crs as ccrs
+
+# types
+BBox = namedtuple("BBox", ["north", "south", "east", "west"])  # type: ignore
+
+
+# coordinate function
+def to_mercator(bbox: BBox) -> BBox:
+    proj = ccrs.Mercator()
+    west, south = proj.transform_point(bbox.west, bbox.south, ccrs.PlateCarree())
+    east, north = proj.transform_point(bbox.east, bbox.north, ccrs.PlateCarree())
+    return BBox(north=north, south=south, east=east, west=west)
+
+
+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)

parabellum/env.py

@@ -2,15 +2,14 @@
 
 import jax.numpy as jnp
 import jax
-import numpy as np
-from jax import random, Array
-from jax import jit
+from jax import random, Array, vmap, jit
 from flax.struct import dataclass
 import chex
-from jax import vmap
 from jaxmarl.environments.smax.smax_env import SMAX
+
 from typing import Tuple, Dict, cast
 from functools import partial
+from parabellum import tps, geo
 
 
 @dataclass
@@ -18,8 +17,8 @@ class Scenario:
     """Parabellum scenario"""
 
     place: str
-    terrain_raster: jnp.ndarray
-    unit_starting_sectors: jnp.ndarray
+    terrain_raster: 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
@@ -27,9 +26,11 @@ class Scenario:
     smacv2_position_generation: bool = False
     smacv2_unit_type_generation: bool = False
 
+
 @dataclass
 class State:
-    unit_positions: Array
+    # terrain: Array
+    unit_positions: Array  # fsfds
     unit_alive: Array
     unit_teams: Array
     unit_health: Array
@@ -41,71 +42,89 @@ class State:
     terminal: bool
 
 
-# default scenario
-scenarios = {
-    "default": Scenario(
-        "Identity Town",
-        jnp.eye(64, dtype=jnp.uint8),
-        jnp.array([[0, 0, 0.2, 0.2], [0.7,0.7,0.2,0.2]]),
-        jnp.zeros((19,), dtype=jnp.uint8),
-        9,
-        10,
-    )
-}
-
 
-def make_scenario(place, terrain_raster, unit_starting_sectors, allies_type, n_allies, enemies_type, n_enemies):
+def make_scenario(
+    place,
+    size,
+    unit_starting_sectors,
+    allies_type,
+    n_allies,
+    enemies_type,
+    n_enemies,
+):
+    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)
+        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)
+        assert len(enemies_type) == n_enemies
         enemies = enemies_type
     unit_types = jnp.array(allies + enemies, dtype=jnp.uint8)
-    return Scenario(place, terrain_raster, unit_starting_sectors, unit_types, n_allies, n_enemies)
+    return Scenario(
+        place, terrain, unit_starting_sectors, unit_types, n_allies, n_enemies  # type: ignore
+    )
 
 
-def spawn_fn(pool, offset: jnp.ndarray, n: int, rng: jnp.ndarray):
+def spawn_fn(rng: jnp.ndarray, units_spawning_sectors):
     """Spawns n agents on a map."""
-    rng, key_start, key_noise = random.split(rng, 3)
-    noise = random.uniform(key_noise, (n, 2)) * 0.5
-
-    # select n random (x, y)-coords where sector == True
-    idxs = random.choice(key_start, pool[0].shape[0], (n,), replace=False)
-    coords = jnp.array([pool[0][idxs], pool[1][idxs]]).T
-
-    return coords + noise + offset
-
-
-def sector_fn(terrain: jnp.ndarray, sector_id: int):
-    """return sector slice of terrain"""
-    width, height = terrain.shape
-    coordx, coordy = sector_id // 5 * width // 5, sector_id % 5 * height // 5
-    sector = terrain[coordx : coordx + width // 5, coordy : coordy + height // 5] == 0
-    offset = jnp.array([coordx, coordy])
-    # sector is jnp.nonzero
-    return jnp.nonzero(sector), offset
-
-
-def sector_fn(terrain: jnp.ndarray, sector: jnp.ndarray):
-    """return sector slice of terrain"""
-    width, height = terrain.shape
-    coordx, coordy = int(sector[0] * width), int(sector[1] * height)
-    sector = terrain[coordy : coordy + int(sector[3] * height), coordx : coordx + int(sector[2] * width)] == 0
-    offset = jnp.array([coordx, coordy])
-    # sector is jnp.nonzero
-    return jnp.nonzero(sector.T), offset
+    spawn_positions = []
+    for sector in units_spawning_sectors:
+        rng, key_start, key_noise = random.split(rng, 3)
+        noise = 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[coordy : coordy + int(sector[3] * height), coordx : coordx + int(sector[2] * width)] == 0)
+        valid = jnp.nonzero(sector.T)
+        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_raster.shape
+        map_height, map_width = scenario.terrain_raster.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_raster = scenario.terrain_raster
         self.unit_starting_sectors = scenario.unit_starting_sectors
@@ -113,21 +132,18 @@ class Environment(SMAX):
         # 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
+        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.team0_sector, self.team0_sector_offset = sector_fn(self.terrain_raster, self.unit_starting_sectors[0]) # sector_fn(self.terrain_raster, 0)
-        self.team1_sector, self.team1_sector_offset = sector_fn(self.terrain_raster, self.unit_starting_sectors[1]) # sector_fn(self.terrain_raster, 24)
+        self._push_units_away = lambda state, firmness=1: state  # overwrite push units
+        self.spawning_sectors = sectors_fn(self.unit_starting_sectors, scenario.terrain_raster.building + scenario.terrain_raster.water)
+        self.resolution = self.terrain_raster.building.shape[0] + self.terrain_raster.building.shape[1]
+        self.t = jnp.tile(jnp.linspace(0, 1, self.resolution), (2, self.resolution))
 
 
-    @partial(jax.jit, static_argnums=(0,))
-    def reset(self, rng: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]:
+    def reset(self, rng: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]: # type: ignore
         """Environment-specific reset."""
-        ally_key, enemy_key = jax.random.split(rng)
-        team_0_start = spawn_fn(self.team0_sector, self.team0_sector_offset, self.num_allies, ally_key)
-        team_1_start = spawn_fn(self.team1_sector, self.team1_sector_offset, self.num_enemies, enemy_key)
-        unit_positions = jnp.concatenate([team_0_start, team_1_start])
+        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,))
@@ -145,13 +161,66 @@ class Environment(SMAX):
             time=0,
             terminal=False,
             unit_weapon_cooldowns=unit_weapon_cooldowns,
+            # terrain=self.terrain_raster,
         )
-        state = self._push_units_away(state)  # type: ignore
+        state = self._push_units_away(state)  # type: ignore  could be slow
         obs = self.get_obs(state)
         world_state = self.get_world_state(state)
-        obs["world_state"] = jax.lax.stop_gradient(world_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
+        obs.pop("world_state")
+        if not self.reset_when_done:
+            for key in dones.keys():
+                dones[key] = False
+        return obs, state, rewards, dones, infos
+
+    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_raster.building + self.terrain_raster.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
@@ -171,7 +240,19 @@ class Environment(SMAX):
             pos
             + firmness * jnp.sum(delta_matrix * overlap_term[:, :, None], axis=1) / 2
         )
-        return unit_positions
+        return jnp.where(self.unit_type_pushable[unit_types][:, None], unit_positions, pos)
+
+    def has_line_of_sight(self, source, target, raster_input):  # this is tooooo slow TODO: make it fast
+        # we could compute this for units in sight only using a switch
+
+        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[1, :], cells[0, :]].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
@@ -180,25 +261,17 @@ class Environment(SMAX):
         state: State,
         actions: Tuple[chex.Array, chex.Array],
     ) -> State:
-        @partial(jax.vmap, in_axes=(None, None, 0, 0))
-        def intersect_fn(pos, new_pos, obs, obs_end):
-            d1 = jnp.cross(obs - pos, new_pos - pos)
-            d2 = jnp.cross(obs_end - pos, new_pos - pos)
-            d3 = jnp.cross(pos - obs, obs_end - obs)
-            d4 = jnp.cross(new_pos - obs, obs_end - obs)
-            return (d1 * d2 <= 0) & (d3 * d4 <= 0)
-
-        def raster_crossing(pos, new_pos):
+        def raster_crossing(pos, new_pos, mask: jnp.ndarray):
             pos, new_pos = pos.astype(jnp.int32), new_pos.astype(jnp.int32)
-            raster = jnp.copy(self.terrain_raster)
             minimum = jnp.minimum(pos, new_pos)
             maximum = jnp.maximum(pos, new_pos)
-            raster = jnp.where(jnp.arange(raster.shape[0]) >= minimum[0], raster, 0)
-            raster = jnp.where(jnp.arange(raster.shape[0]) <= maximum[0], raster, 0)
-            raster = jnp.where(jnp.arange(raster.shape[1]) >= minimum[1], raster.T, 0).T
-            raster = jnp.where(jnp.arange(raster.shape[1]) <= maximum[1], raster.T, 0).T
-            return jnp.any(raster)
-        
+            mask = jnp.where(jnp.arange(mask.shape[0]) >= minimum[0], mask, 0)
+            mask = jnp.where(jnp.arange(mask.shape[0]) <= maximum[0], mask, 0)
+            mask = jnp.where(jnp.arange(mask.shape[1]) >= minimum[1], mask.T, 0).T
+            mask = jnp.where(jnp.arange(mask.shape[1]) <= maximum[1], mask.T, 0).T
+            return jnp.any(mask)
+
+
         def update_position(idx, vec):
             # Compute the movements slightly strangely.
             # The velocities below are for diagonal directions
@@ -214,13 +287,13 @@ class Environment(SMAX):
             )
             # avoid going out of bounds
             new_pos = jnp.maximum(
-                jnp.minimum(new_pos, jnp.array([self.map_width, self.map_height])),
+                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)
+            clash = raster_crossing(pos, new_pos, self.terrain_raster.building + self.terrain_raster.water)
             new_pos = jnp.where(clash, pos, new_pos)
 
             #######################################################################
@@ -285,7 +358,7 @@ class Environment(SMAX):
             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.bool_)  # type: ignore
                 .astype(jnp.float32)
             )
             bystander_health_diff = (
@@ -333,11 +406,14 @@ class Environment(SMAX):
 
         # units push each other
         new_pos = self._our_push_units_away(pos, state.unit_types)
-        clash = jax.vmap(raster_crossing)(pos, new_pos)
+        clash = jax.vmap(raster_crossing, in_axes=(0, 0, None))(pos, new_pos, self.terrain_raster.building + self.terrain_raster.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, self.map_height])),jnp.zeros((2,)),)
-        
+        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
@@ -380,26 +456,39 @@ class Environment(SMAX):
         )
         return state
 
-
 if __name__ == "__main__":
     n_envs = 4
 
-    env = Environment(scenarios["default"])
+
+    n_allies = 10
+    scenario_kwargs = {"allies_type": 0, "n_allies": n_allies, "enemies_type": 0, "n_enemies": n_allies,
+                        "place": "Vesterbro, Copenhagen, Denmark", "size": 256, "unit_starting_sectors":
+                            [([i for i in range(n_allies)], [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 = []
 
-    print(state.unit_positions)
-    exit()
 
-    for i in range(10):
+    from tqdm import tqdm
+    import time
+    step = vmap(jit(env.step))
+    tic = time.time()
+    for i in tqdm(range(10)):
         rng, act_rng, step_rng = random.split(rng, 3)
         act_key = random.split(act_rng, (len(env.agents), n_envs))
+        print(tic - time.time())
         act = {
             a: vmap(env.action_space(a).sample)(act_key[i])
             for i, a in enumerate(env.agents)
         }
+        print(tic - time.time())
         step_key = random.split(step_rng, n_envs)
+        print(tic - time.time())
         state_seq.append((step_key, state, act))
-        obs, state, reward, done, infos = vmap(env.step)(step_key, state, act)
+        print(tic - time.time())
+        obs, state, reward, done, infos = step(step_key, state, act)
+        print(tic - time.time())
+        tic = time.time()

parabellum/geo.py

@@ -0,0 +1,100 @@
+# %% geo.py
+#   script for geospatial level generation
+# by: Noah Syrkis
+
+# %% Imports
+from parabellum import tps
+import rasterio
+from rasterio import features, transform
+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
+
+# %% Types
+Coords = Tuple[float, float]
+BBox = namedtuple("BBox", ["north", "south", "east", "west"])  # type: ignore
+
+# %% Constants
+provider = cx.providers.Stadia.StamenTerrain(  # type: ignore
+    api_key="86d0d32b-d2fe-49af-8db8-f7751f58e83f"
+)
+provider["url"] = provider["url"] + "?api_key={api_key}"
+tags = {"building": True, "water": True, "landuse": "forest"}  #  "road": True}
+
+
+# %% Coordinate function
+def get_coordinates(place: str) -> Coords:
+    geolocator = Nominatim(user_agent="parabellum")
+    point = geolocator.geocode(place)
+    return point.latitude, point.longitude  # type: ignore
+
+
+def get_bbox(place: str, buffer) -> BBox:
+    """Get bounding box from place name in crs 4326."""
+    coords = get_coordinates(place)
+    north = distance(meters=buffer).destination(coords, bearing=0).latitude
+    south = distance(meters=buffer).destination(coords, bearing=180).latitude
+    east = distance(meters=buffer).destination(coords, bearing=90).longitude
+    west = distance(meters=buffer).destination(coords, bearing=270).longitude
+    return BBox(north, south, east, west)
+
+
+def basemap_fn(bbox: BBox, gdf) -> Array:
+    fig, ax = plt.subplots(figsize=(20, 20), subplot_kw={"projection": ccrs.Mercator()})
+    gdf.plot(ax=ax, color="black", alpha=0, edgecolor="black")  # type: ignore
+    cx.add_basemap(ax, crs=gdf.crs, source=provider, zoom="auto") # type: ignore
+    bbox = gdf.total_bounds
+    ax.set_extent([bbox[0], bbox[2], bbox[1], bbox[3]], crs=ccrs.Mercator())  # type: ignore
+    plt.axis("off")
+    plt.tight_layout()
+    fig.canvas.draw()
+    image = jnp.array(fig.canvas.renderer._renderer)  # type: ignore
+    plt.close(fig)
+    return image
+
+
+def geography_fn(place, buffer):
+    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")
+    raster = raster_fn(gdf, shape=(buffer, buffer))
+    basemap = basemap_fn(bbox, gdf)
+    terrain = tps.Terrain(building=raster[0], water=raster[1], forest=raster[2], basemap=basemap)
+    return terrain
+
+
+def raster_fn(gdf, shape) -> Array:
+    bbox = gdf.total_bounds
+    t = transform.from_bounds(*bbox, *shape)  # type: ignore
+    raster = jnp.array([feature_fn(t, feature, gdf, shape) for feature in ["building", "water", "landuse"]])
+    return raster
+
+def feature_fn(t, feature, gdf, shape):
+    if feature not in gdf.columns:
+        return jnp.zeros(shape)
+    gdf = gdf[~gdf[feature].isna()]
+    raster = features.rasterize(gdf.geometry, out_shape=shape, transform=t, fill=0)  # type: ignore
+    return raster
+
+place = "Thun, Switzerland"
+terrain = geography_fn(place, 800)
+# %%
+fig, axes = plt.subplots(1, 5, figsize=(20, 20))
+axes[0].imshow(terrain.building, cmap="gray")
+axes[1].imshow(terrain.water, cmap="gray")
+axes[2].imshow(terrain.forest, cmap="gray")
+axes[3].imshow(terrain.building + terrain.water + terrain.forest)
+axes[4].imshow(terrain.basemap)

parabellum/map.py

@@ -1,100 +1,95 @@
-# map.py
-    # parabellum map functions
+# ludens.py
+#    script for fucking around and finding out
 # by: Noah Syrkis
 
-# imports
-import jax.numpy as jnp
-from geopy.geocoders import Nominatim
-import geopandas as gpd
+
+# %% Imports
+# import parabellum as pb
+import matplotlib.pyplot as plt
 import osmnx as ox
+from geopy.geocoders import Nominatim
+import numpy as np
 import contextily as cx
-import matplotlib.pyplot as plt
+import jax.numpy as jnp
+import geopandas as gpd
+import rasterio
 from rasterio import features
-import rasterio.transform
-from typing import Optional, Tuple
-from geopy.location import Location
 from shapely.geometry import Point
-import os
-import pickle 
+from typing import List
 
-# constants
+# %% Constants
 geolocator = Nominatim(user_agent="parabellum")
-BUILDING_TAGS = {"building": True}
-
-def get_location(place: str) -> Tuple[float, float]:
-    """Get coordinates for a given place."""
-    coords: Optional[Location] = geolocator.geocode(place)  # type: ignore
-    if coords is None:
-        raise ValueError(f"Could not geocode the place: {place}")
-    return (coords.latitude, coords.longitude)
-
-def get_building_geometry(point: Tuple[float, float], size: int) -> gpd.GeoDataFrame:
-    """Get building geometry for a given point and size."""
-    geometry = ox.features_from_point(point, tags=BUILDING_TAGS, dist=size // 2)
-    return gpd.GeoDataFrame(geometry).set_crs("EPSG:4326")
-
-def rasterize_geometry(gdf: gpd.GeoDataFrame, size: int) -> jnp.ndarray:
-    """Rasterize geometry and return as a JAX array."""
-    w, s, e, n = gdf.total_bounds
-    transform = rasterio.transform.from_bounds(w, s, e, n, size, size)
-    raster = features.rasterize(gdf.geometry, out_shape=(size, size), transform=transform)
-    return jnp.array(jnp.flip(raster, 0) ).astype(jnp.uint8)
-
-# +
-def get_from_cache(place, size):
-    if os.path.exists("./cache"):
-        name = str(hash((place, size))) + ".pk"
-        if os.path.exists("./cache/" + name):
-            with open("./cache/" + name, "rb") as f:
-                (mask, base) = pickle.load(f)
-            return (mask, base.astype(jnp.int64))
-    return (None, None)
-    
-def save_in_cache(place, size, mask, base):
-    if not os.path.exists("./cache"):
-        os.makedirs("./cache")
-    name = str(hash((place, size))) + ".pk"
-    with open("./cache/" + name, "wb") as f:
-        pickle.dump((mask, base), f)
-
-def terrain_fn(place: str, size: int = 1000, with_cache: bool = True) -> Tuple[jnp.ndarray, jnp.ndarray]:
-    """Returns a rasterized map of buildings for a given location."""
-    if with_cache:
-        mask, base = get_from_cache(place, size)
-    if not with_cache or mask is None:
-        point = get_location(place)
-        gdf = get_building_geometry(point, size)
-        mask = rasterize_geometry(gdf, size)
-        base = get_basemap(place, size)
-        if with_cache:
-            save_in_cache(place, size, mask, base)
-    return mask, base
-
-
-# -
-
-def get_basemap(place: str, size: int = 1000) -> jnp.ndarray:
-    """Returns a basemap for a given place as a JAX array."""
-    point = get_location(place)
-    gdf = get_building_geometry(point, size)
-    basemap, _ = cx.bounds2img(*gdf.total_bounds, ll=True)
-    # get the middle size x size square
-    basemap = basemap[(basemap.shape[0] - size) // 2:(basemap.shape[0] + size) // 2,
-                        (basemap.shape[1] - size) // 2:(basemap.shape[1] + size) // 2]
-    return basemap # jnp.array(jnp.rot90(basemap, 2)).astype(jnp.uint8)
-
-
-if __name__ == "__main__":
-    place = "Cauvicourt, 14190, France"
-    mask, base = terrain_fn(place, 500)
-    
-    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
-    ax[0].imshow(jnp.flip(mask,0)) # type: ignore
-    ax[1].imshow(base) # type: ignore
-    ax[2].imshow(base) # type: ignore
-    ax[2].imshow(jnp.flip(mask,0), alpha=jnp.flip(mask,0)) # type: ignore
-    plt.show()
-
-
-
-
+source = cx.providers.OpenStreetMap.Mapnik  # type: ignore
+
+
+def get_raster(
+    place: str, meters: int = 1000, tags: List[dict] | dict = {"building": True}
+) -> jnp.ndarray:
+    # look here for tags https://wiki.openstreetmap.org/wiki/Map_features
+    def aux(place, tag):
+        """Rasterize geometry and return as a JAX array."""
+        place = geolocator.geocode(place)  # type: ignore
+        point = place.latitude, place.longitude  # type: ignore  # confusing order of lat/lon
+        geom = ox.features_from_point(point, tags=tag, dist=meters // 2)
+        gdf = gpd.GeoDataFrame(geom).set_crs("EPSG:4326")
+        # crop everythin outside of the meters x meters square
+        gdf = gdf.cx[
+            place.longitude - meters / 2 : place.longitude + meters / 2,
+            place.latitude - meters / 2 : place.latitude + meters / 2,
+        ]
+
+        # bounds should be meters, meters
+        t = rasterio.transform.from_bounds(*bounds, meters, meters)  # type: ignore
+        raster = features.rasterize(
+            gdf.geometry, out_shape=(meters, meters), transform=t
+        )
+        return jnp.array(raster)
+
+    if isinstance(tags, dict):
+        return aux(place, tags)
+    else:
+        return jnp.stack([aux(place, tag) for tag in tags])
+
+
+def get_basemap(
+    place: str, size: int = 1000
+) -> np.ndarray:  # TODO: image is slightly off from raster. Fix this.
+    # Create a GeoDataFrame with the center point
+    place = geolocator.geocode(place)  # type: ignore
+    lon, lat = place.longitude, place.latitude  # type: ignore
+    gdf = gpd.GeoDataFrame(geometry=[Point(lon, lat)], crs="EPSG:4326")
+    gdf = gdf.to_crs("EPSG:3857")
+
+    # Create a buffer around the center point
+    # buffer = gdf.buffer(size)  # type: ignore
+    buffer = gdf
+    bounds = buffer.total_bounds  # i think this is wrong, since it ignores empty space
+    # modify bounds to include empty space
+    bounds = (bounds[0] - size, bounds[1] - size, bounds[2] + size, bounds[3] + size)
+
+    # Create a figure and axis
+    dpi = 300
+    fig, ax = plt.subplots(figsize=(size / dpi, size / dpi), dpi=dpi)
+    buffer.plot(ax=ax, facecolor="none", edgecolor="red", linewidth=0)
+
+    # Calculate the zoom level for the basemap
+
+    # Add the basemap to the axis
+    cx.add_basemap(ax, source=source, zoom="auto", attribution=False)
+
+    # Set the x and y limits of the axis
+    ax.set_xlim(bounds[0], bounds[2])
+    ax.set_ylim(bounds[1], bounds[3])
+
+    # convert the image (without axis or border) to a numpy array
+    plt.axis("off")
+    plt.tight_layout()
+
+    # remove whitespace
+    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
+    fig.canvas.draw()
+
+    image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)  # type: ignore
+    image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return jnp.array(image)  # type: ignore

parabellum/pcg.py

@@ -0,0 +1,53 @@
+# pcg.py
+#   procedural content generation
+# by: Noah Syrkis
+
+# %% Imports
+from jax import random, vmap
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+from functools import partial
+
+
+# %% Functions
+seed = 0
+n = 100
+rng = random.PRNGKey(seed)
+Y = random.uniform(rng, (n,))
+
+
+def g(t):
+    return (1 - jnp.cos(jnp.pi * t)) / 2
+
+
+def lerp(a, b, t):
+    t -= jnp.floor(t)  # the fractional part of t
+    return (1 - t) * a + t * b
+
+
+def cerp(a, b, t):
+    t -= jnp.floor(t)  # the fractional part of t
+    return g(1 - t) * a + g(t) * b
+
+
+def body_fn(x):
+    i = jnp.floor(x).astype(jnp.uint8)
+    return cerp(Y[i], Y[i + 1], x)
+
+
+@partial(vmap, in_axes=(None, 0, None))
+def noise_fn(y, t, n):
+    return y[t % n]
+
+
+@vmap
+def perlin_fn(t):
+    return noise_fn(Y, t * jnp.arange(n * 3), n)
+
+
+xs = jnp.linspace(0, 1, 1000)
+noise = perlin_fn(2 ** jnp.arange(3)).sum(0)
+
+fig, ax = plt.subplots(figsize=(20, 4), dpi=100)
+ax.set_ylim(0, 1)
+ax.plot(noise / noise.max())

parabellum/run.py

@@ -21,7 +21,7 @@ bg = (0, 0, 0) if darkdetect.isDark() else (255, 255, 255)
 
 
 # types
-State = jaxmarl.environments.smax.smax_env.State # type: ignore
+State = jaxmarl.environments.smax.smax_env.State  # type: ignore
 Obs = Reward = Done = Action = Dict[str, jnp.ndarray]
 StateSeq = List[Tuple[jnp.ndarray, State, Action]]
 

parabellum/tps.py

@@ -0,0 +1,16 @@
+# 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/vis.py

@@ -3,26 +3,19 @@ Visualizer for the Parabellum environment
 """
 
 # Standard library imports
-from functools import partial
-from typing import Optional, List, Tuple
+from typing import Optional, Tuple
 import cv2
-from PIL import Image
 
 # JAX and JAX-related imports
 import jax
 from chex import dataclass
-import chex
-from jax import vmap, tree_util, Array, jit
+from jax import vmap, Array
 import jax.numpy as jnp
-from jaxmarl.environments.multi_agent_env import MultiAgentEnv
-from jaxmarl.environments.smax import SMAX
 from jaxmarl.viz.visualizer import SMAXVisualizer
 
 # Third-party imports
 import numpy as np
 import pygame
-import cv2
-from tqdm import tqdm
 
 # Local imports
 import parabellum as pb
@@ -35,12 +28,12 @@ class Skin:
     maskmap: Array  # maskmap of buildings
     bg: Tuple[int, int, int] = (255, 255, 255)
     fg: Tuple[int, int, int] = (0, 0, 0)
-    ally: Tuple[int, int, int]  = (0, 255, 0)
-    enemy: Tuple[int, int, int]  = (255, 0, 0)
-    pad: int  = 100
-    size: int  = 1000  # excluding padding
+    ally: Tuple[int, int, int] = (0, 255, 0)
+    enemy: Tuple[int, int, int] = (255, 0, 0)
+    pad: int = 100
+    size: int = 1000  # excluding padding
     fps: int = 24
-    vis_size: int = 1000   # size of the map in Vis (exluding padding)
+    vis_size: int = 1000  # size of the map in Vis (exluding padding)
     scale: Optional[float] = None
 
 
@@ -57,10 +50,23 @@ class Visualizer(SMAXVisualizer):
         self.env = env
 
     def animate(self, save_fname: Optional[str] = "output/parabellum", view=None):
-        expanded_state_seq, expanded_action_seq = expand_fn(self.env, self.state_seq, self.action_seq)
-        state_seq_seq, action_seq_seq = unbatch_fn(expanded_state_seq, expanded_action_seq)
-        for idx, (state_seq, action_seq) in enumerate(zip(state_seq_seq, action_seq_seq)):
-            animate_fn(self.env, self.skin, self.image, state_seq, action_seq, f"{save_fname}_{idx}.mp4")
+        expanded_state_seq, expanded_action_seq = expand_fn(
+            self.env, self.state_seq, self.action_seq
+        )
+        state_seq_seq, action_seq_seq = unbatch_fn(
+            expanded_state_seq, expanded_action_seq
+        )
+        for idx, (state_seq, action_seq) in enumerate(
+            zip(state_seq_seq, action_seq_seq)
+        ):
+            animate_fn(
+                self.env,
+                self.skin,
+                self.image,
+                state_seq,
+                action_seq,
+                f"{save_fname}_{idx}.mp4",
+            )
 
 
 # functions
@@ -70,22 +76,29 @@ def animate_fn(env, skin, image, state_seq, action_seq, save_fname):
     for idx, (state_tup, action) in enumerate(zip(state_seq, action_seq)):
         frames += [frame_fn(env, skin, image, state_tup[1], action, idx)]
     # use cv2 to write frames to video
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # type: ignore
-    out = cv2.VideoWriter(save_fname, fourcc, skin.fps, (skin.size + skin.pad * 2, skin.size + skin.pad * 2))
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # type: ignore
+    out = cv2.VideoWriter(
+        save_fname,
+        fourcc,
+        skin.fps,
+        (skin.size + skin.pad * 2, skin.size + skin.pad * 2),
+    )
     for frame in frames:
         out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
     out.release()
     pygame.quit()
 
 
-def init_frame(env, skin, image, state: pb.State, action: Array, idx: int) -> pygame.Surface:
+def init_frame(
+    env, skin, image, state: pb.State, action: Array, idx: int
+) -> pygame.Surface:
     dims = (skin.size + skin.pad * 2, skin.size + skin.pad * 2)
     frame = pygame.Surface(dims, pygame.SRCALPHA | pygame.HWSURFACE)
     return frame
 
 
 def transform_frame(env, skin, frame):
-    #frame = np.rot90(pygame.surfarray.pixels3d(frame).swapaxes(0, 1), 2)
+    # frame = np.rot90(pygame.surfarray.pixels3d(frame).swapaxes(0, 1), 2)
     frame = np.flip(pygame.surfarray.pixels3d(frame).swapaxes(0, 1), 0)
     return frame
 
@@ -102,7 +115,7 @@ def frame_fn(env, skin, image, state: pb.State, action: Array, idx: int) -> np.n
 
 
 def render_background(env, skin, image, frame, state, action):
-    coords = (skin.pad-5, skin.pad-5, skin.size+10, skin.size+10)
+    coords = (skin.pad - 5, skin.pad - 5, skin.size + 10, skin.size + 10)
     frame.fill(skin.bg)
     frame.blit(image, coords)
     pygame.draw.rect(frame, skin.fg, coords, 3)
@@ -116,6 +129,7 @@ def render_action(env, skin, image, frame, state, action):
 def render_bullet(env, skin, image, frame, state, action):
     return frame
 
+
 def render_agents(env, skin, image, frame, state, action):
     units = state.unit_positions, state.unit_teams, state.unit_types, state.unit_health
     for idx, (pos, team, kind, health) in enumerate(zip(*units)):
@@ -136,7 +150,11 @@ def text_fn(text):
 
 def image_fn(skin: Skin):  # TODO:
     """Create an image for background (basemap or maskmap)"""
-    motif = cv2.resize(np.array(skin.maskmap.T), (skin.size, skin.size), interpolation=cv2.INTER_NEAREST).astype(np.uint8)
+    motif = cv2.resize(
+        np.array(skin.maskmap.T),
+        (skin.size, skin.size),
+        interpolation=cv2.INTER_NEAREST,
+    ).astype(np.uint8)
     motif = (motif > 0).astype(np.uint8)
     image = np.zeros((skin.size, skin.size, 3), dtype=np.uint8) + skin.bg
     image[motif == 1] = skin.fg
@@ -150,7 +168,9 @@ def unbatch_fn(state_seq, action_seq):
     if is_multi_run(state_seq):
         n_envs = state_seq[0][1].unit_positions.shape[0]
         state_seq_seq = [jax.tree_map(lambda x: x[i], state_seq) for i in range(n_envs)]
-        action_seq_seq = [jax.tree_map(lambda x: x[i], action_seq) for i in range(n_envs)]
+        action_seq_seq = [
+            jax.tree_map(lambda x: x[i], action_seq) for i in range(n_envs)
+        ]
     else:
         state_seq_seq = [state_seq]
         action_seq_seq = [action_seq]
@@ -161,7 +181,9 @@ def expand_fn(env, state_seq, action_seq):
     """Expand the state sequence"""
     fn = env.expand_state_seq
     state_seq = vmap(fn)(state_seq) if is_multi_run(state_seq) else fn(state_seq)
-    action_seq = [action_seq[i // env.world_steps_per_env_step] for i in range(len(state_seq))]
+    action_seq = [
+        action_seq[i // env.world_steps_per_env_step] for i in range(len(state_seq))
+    ]
     return state_seq, action_seq
 
 

{parabellum-0.2.25.dist-info → parabellum-0.3.0.dist-info}/METADATA

@@ -1,35 +1,37 @@
 Metadata-Version: 2.1
 Name: parabellum
-Version: 0.2.25
+Version: 0.3.0
 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,<4.0
+Requires-Python: >=3.11,<3.12
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
+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: einops (>=0.8.0,<0.9.0)
 Requires-Dist: folium (>=0.17.0,<0.18.0)
-Requires-Dist: geopandas (>=1.0.0,<2.0.0)
 Requires-Dist: geopy (>=2.4.1,<3.0.0)
 Requires-Dist: ipykernel (>=6.29.5,<7.0.0)
 Requires-Dist: jax (==0.4.17)
 Requires-Dist: jaxmarl (==0.0.3)
+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: osmnx (>=1.9.3,<2.0.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
 Description-Content-Type: text/markdown

parabellum-0.3.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+parabellum/__init__.py,sha256=vqQbvsTT_zcLThZ7fLoJ6cMAZbEeGIJDFyCkHmovfOY,392
+parabellum/aid.py,sha256=BPabjN4BUq1HRhkwbc9pCNsXSF_ALiG8W8cHWTWeEH4,900
+parabellum/env.py,sha256=VV3VK7TTkianihqJopRbY0vlRWOquu-VTrc9ep0PSTk,21304
+parabellum/geo.py,sha256=xkj6iJqN076tRbaG38Sq7gtwKSNzxI37msRLnpn5JV0,3561
+parabellum/gun.py,sha256=nvsJdcZ2Qd6lbPlAgsUiaLhstTi1UdLQ8kOnbCenucY,2618
+parabellum/map.py,sha256=9AV0PIqInXcWWojzHshy3X42Nm3ZDq0O1NG-6fQ9Wgw,3345
+parabellum/pcg.py,sha256=d8KC_lbc4WUUUPaTdPJSx27VMGioys3jSGOWJ-2EahU,968
+parabellum/run.py,sha256=Q53__AxzROZNgfZLVU5LDdcT61UMCkmQ_Q5wWUIrnqo,3473
+parabellum/tps.py,sha256=3tVqo42ggE8idZn500C0X2pS9TmYndgBzlAG7Yj2Wz8,252
+parabellum/vis.py,sha256=ABHveJj0fLRWkxOv3LFIXK20QtdGhjskuFLsp7iTFu0,6185
+parabellum-0.3.0.dist-info/METADATA,sha256=FugXwz25bAPYKlIfqFc7dGVtPupse5zHYapmqBWopE8,2740
+parabellum-0.3.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+parabellum-0.3.0.dist-info/RECORD,,

parabellum-0.2.25.dist-info/RECORD

@@ -1,10 +0,0 @@
-parabellum/__init__.py,sha256=-5cWXJkHnfH_CbhTEall8Wak8McAFXZHP1L8Fu7Uo5k,373
-parabellum/aid.py,sha256=HWST27inTFXcp8b11izJF0U7N7DZnRTIS3n1Qfa-Ko4,106
-parabellum/env.py,sha256=H1VwxHj8KqbjBJ8b7NMxkAg3Q0qwVzGulrigc26Tzkc,16663
-parabellum/gun.py,sha256=nvsJdcZ2Qd6lbPlAgsUiaLhstTi1UdLQ8kOnbCenucY,2618
-parabellum/map.py,sha256=UwMqwySasX5oLw9v5YMJARAPwvQThLTRW36NpbwvBC8,3564
-parabellum/run.py,sha256=EO_F7VPwayatpSHrcbSahtinsV4QObhcx0jo-4KZO1E,3472
-parabellum/vis.py,sha256=q7_OIMjzt-7nBOojVVW7Wiiq9ojsjaltIsH6eOOxPKk,6116
-parabellum-0.2.25.dist-info/METADATA,sha256=UCuRLYhSUxnebs5pRs1FaWK67PUJw0kgsJ2UaaxNM9Q,2671
-parabellum-0.2.25.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-parabellum-0.2.25.dist-info/RECORD,,