parabellum 0.2.26__py3-none-any.whl → 0.3.1__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 terrain_db
+from . import env
+from . import tps
 # from .run import run
 
 __all__ = [
+    "env",
+    "terrain_db",
+    "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,16 @@
 
 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 math import ceil
+
 from typing import Tuple, Dict, cast
 from functools import partial
+from parabellum import tps, geo, terrain_db
 
 
 @dataclass
@@ -18,8 +19,8 @@ class Scenario:
     """Parabellum scenario"""
 
     place: str
-    terrain_raster: jnp.ndarray
-    unit_starting_sectors: jnp.ndarray
+    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
@@ -27,9 +28,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,93 +44,156 @@ 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,
+):
+    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)
+        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)
-
-
-def spawn_fn(pool, offset: jnp.ndarray, n: int, rng: jnp.ndarray):
-    """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
+    return Scenario(
+        place,
+        terrain,
+        unit_starting_sectors,  # type: ignore
+        unit_types,
+        n_allies,
+        n_enemies,
+    )
 
 
-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 scenario_fn(place):
+    # scenario function for Noah, cos the one above is confusing
+    terrain = geo.geography_fn(place)
+    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 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
+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_raster.shape
+        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_raster = scenario.terrain_raster
+        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
+        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.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))
 
-    @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,18 +211,71 @@ class Environment(SMAX):
             time=0,
             terminal=False,
             unit_weapon_cooldowns=unit_weapon_cooldowns,
+            # terrain=self.terrain,
         )
-        state = self._push_units_away(state)  # type: ignore
+        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)
+        obs["world_state"] = jax.lax.stop_gradient(world_state)
         return obs, state
 
-    def step_env(self, rng, state: State, action: Array):
-        obs, state, rewards, dones, infos = super().step_env(rng, state, action)
+        # 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")
-        return obs, state, rewards, dones, infos
+        # 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.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
@@ -176,7 +295,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):
+        # 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
@@ -185,24 +316,15 @@ 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.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.
@@ -219,13 +341,17 @@ 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.building + self.terrain.water
+            )
             new_pos = jnp.where(clash, pos, new_pos)
 
             #######################################################################
@@ -263,14 +389,11 @@ class Environment(SMAX):
             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 = (
-                (
-                    jnp.linalg.norm(
-                        state.unit_positions[idx] - state.unit_positions[attacked_idx]
-                    )
-                    < self.unit_type_attack_ranges[state.unit_types[idx]]
-                )
+                (distance <= self.unit_type_attack_ranges[state.unit_types[idx]])
                 & state.unit_alive[idx]
                 & state.unit_alive[attacked_idx]
             )
@@ -281,21 +404,28 @@ class Environment(SMAX):
                 -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]]
-            )
+            # 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]]
+            # )
 
             #########################################################
             #########################################################
@@ -319,29 +449,34 @@ class Environment(SMAX):
                 health_diff,
                 attacked_idx,
                 cooldown_diff,
-                (bystander_health_diff, bystander_idxs),
+                # (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, (bystander) = (
-                update_agent_health(idx, attack_action, key)
+            health_diff, attacked_idxes, cooldown_diff = update_agent_health(
+                idx, attack_action, key
             )
 
-            return new_pos, (health_diff, attacked_idxes), cooldown_diff, bystander
+            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, bystander = jax.vmap(
+        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)(pos, new_pos)
+        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, 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.
@@ -370,8 +505,8 @@ class Environment(SMAX):
         #########################################################
         ############################ subtracting bystander health
 
-        _, bystander_health_diff = bystander
-        unit_health -= bystander_health_diff.sum(axis=0)  # might be axis=1
+        # _, bystander_health_diff = bystander
+        # unit_health -= bystander_health_diff.sum(axis=0)  # might be axis=1
 
         #########################################################
         #########################################################
@@ -389,15 +524,30 @@ class Environment(SMAX):
 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": 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 = []
 
-    print(state.unit_positions)
-    exit()
+    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))
@@ -407,4 +557,5 @@ if __name__ == "__main__":
         }
         step_key = random.split(step_rng, n_envs)
         state_seq.append((step_key, state, act))
-        obs, state, reward, done, infos = vmap(env.step)(step_key, state, act)
+        obs, state, reward, done, infos = step(step_key, state, act)
+        tic = time.time()

parabellum/geo.py

@@ -0,0 +1,130 @@
+# %% 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
+from jax.scipy.signal import convolve
+
+# %% 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,
+    "highway": True,
+    "landuse": [
+        "grass",
+        "forest",
+        "flowerbed",
+        "greenfield",
+        "village_green",
+        "recreation_ground",
+    ],
+    "leisure": "garden",
+}  #  "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(pad=0)
+    fig.canvas.draw()
+    image = jnp.array(fig.canvas.renderer._renderer)  # type: ignore
+    plt.close(fig)
+    return image
+
+
+def geography_fn(place, buffer=400):
+    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 = 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)
+    terrain = tps.Terrain(
+        building=trans(raster[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,
+    )
+    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 tags])
+    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
+
+
+# %%
+if __name__ == "__main__":
+    place = "Thun, Switzerland"
+    terrain = geography_fn(place, 300)
+
+    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/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/terrain_db.py

@@ -0,0 +1,117 @@
+# %%
+import numpy as np
+import jax.numpy as jnp
+from parabellum import tps
+
+
+# %%
+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)
+    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:
+            raster[x, y] = 1
+    return raster
+
+
+# %%
+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
+    return raster
+
+
+# %%
+building_color = jnp.array([201,199,198, 255])
+water_color = jnp.array([193, 237, 254, 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():
+        raster = np.zeros((size, size))
+        if config is not None:
+            for elem in config:
+                if "line" in elem:
+                    raster = map_raster_from_line(raster, elem["line"], 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)
+    args["basemap"] = basemap
+    return tps.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},
+    "water_park": {
+        'building': [{"line":[0.5, 0.5, 0.5, 0.]}], 
+        "water": [{"rect":[0., 0.8, 0.1, 0.05]}, {"rect": [0.2, 0.8, 0.8, 0.05]}], 
+        "forest": [{"rect": [0., 0., 1., 0.2]}]
+    },
+    "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": []
+    },
+}
+
+# %% [raw]
+#     import matplotlib.pyplot as plt
+#     size = 50
+#     raster = np.zeros((size, size))
+#     rect = [0.2, 0.3, 0.05, 0.4]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     rect = [0.4, 0.3, 0.05, 0.4]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     rect = [0.2, 0.3, 0.25, 0.05]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     rect = [0.2, 0.7, 0.25, 0.05]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     rect = [0.6, 0.3, 0.4, 0.45]
+#     raster = map_raster_from_rect(raster, rect, size)
+#     plt.imshow(jnp.rot90(raster))
+
+# %% [markdown]
+# # Main
+
+# %%
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+
+    # %%
+    terrain = make_terrain(db["u_shape"], size=50)
+
+    # %%
+    plt.imshow(jnp.rot90(terrain.basemap))
+
+# %%
+
+# %%

parabellum/tps.py

@@ -0,0 +1,17 @@
+# %%
+# 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.26.dist-info → parabellum-0.3.1.dist-info}/METADATA

@@ -1,42 +1,44 @@
 Metadata-Version: 2.1
 Name: parabellum
-Version: 0.2.26
+Version: 0.3.1
 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
 
-Ultra-scalable JaxMARL based warfare simulation engine developed with Armasuisse funding.
+Ultra-scalable JaxMARL based warfare simulation engine.
 
 [![Documentation Status](https://readthedocs.org/projects/parabellum/badge/?version=latest)](https://parabellum.readthedocs.io/en/latest/?badge=latest)
 
@@ -93,3 +95,4 @@ Full documentation: [parabellum.readthedocs.io](https://parabellum.readthedocs.i
 ## License
 
 MIT
+

parabellum-0.3.1.dist-info/RECORD

@@ -0,0 +1,13 @@
+parabellum/__init__.py,sha256=hIOLir7wgaf_HU4j8uos7PaCrofqPQcr3FcMlBsZyr8,406
+parabellum/aid.py,sha256=BPabjN4BUq1HRhkwbc9pCNsXSF_ALiG8W8cHWTWeEH4,900
+parabellum/env.py,sha256=0mDqQ7-OI-oufBMMBoUt72Kf5OvHr9thilLGzszlICY,22569
+parabellum/geo.py,sha256=PwEwspOppTPrHIXDZB_nGPTnVFIvDzbh2WtqzVKMUaM,4198
+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=XTKlpLAi3ZwoVw4-KS-Eh15NKsBKP-yt8v6FJGUtwdM,3960
+parabellum/tps.py,sha256=of-RBdelAbNCHQZd1I22RWmZkwUEh6f161mx0X_G2tE,257
+parabellum/vis.py,sha256=ABHveJj0fLRWkxOv3LFIXK20QtdGhjskuFLsp7iTFu0,6185
+parabellum-0.3.1.dist-info/METADATA,sha256=RrSY6CrhwpVlbdJzacX2iVkh_MEgtZkZZCPHJWJJjqo,2707
+parabellum-0.3.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+parabellum-0.3.1.dist-info/RECORD,,

parabellum/map.py

@@ -1,100 +0,0 @@
-# map.py
-    # parabellum map functions
-# by: Noah Syrkis
-
-# imports
-import jax.numpy as jnp
-from geopy.geocoders import Nominatim
-import geopandas as gpd
-import osmnx as ox
-import contextily as cx
-import matplotlib.pyplot as plt
-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 
-
-# 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()
-
-
-
-

parabellum-0.2.26.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=H1YGAtUYNJd8OHnZ3sOEXbag5L0WjtJHBGL8ymGPvoE,16898
-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.26.dist-info/METADATA,sha256=AJwdmHRRPG2MosgufeWqQH7LsDRtTvFFGMh1azey9zA,2671
-parabellum-0.2.26.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-parabellum-0.2.26.dist-info/RECORD,,