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

parabellum/__init__.py

@@ -2,14 +2,14 @@ from .env import Environment, Scenario, make_scenario, State
 from .vis import Visualizer, Skin
 from .gun import bullet_fn
 from . import vis
-from . import map
+from . import terrain_db
 from . import env
 from . import tps
 # from .run import run
 
 __all__ = [
     "env",
-    "map",
+    "terrain_db",
     "vis",
     "tps",
     "Environment",

parabellum/env.py

@@ -7,9 +7,11 @@ from flax.struct import dataclass
 import chex
 from jaxmarl.environments.smax.smax_env import SMAX
 
+from math import ceil
+
 from typing import Tuple, Dict, cast
 from functools import partial
-from parabellum import tps, geo
+from parabellum import tps, geo, terrain_db
 
 
 @dataclass
@@ -17,7 +19,7 @@ class Scenario:
     """Parabellum scenario"""
 
     place: str
-    terrain_raster: tps.Terrain
+    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
@@ -42,7 +44,6 @@ class State:
     terminal: bool
 
 
-
 def make_scenario(
     place,
     size,
@@ -52,19 +53,27 @@ def make_scenario(
     enemies_type,
     n_enemies,
 ):
-    terrain = geo.geography_fn(place, size)
+    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
+        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:
@@ -78,7 +87,30 @@ def make_scenario(
         enemies = enemies_type
     unit_types = jnp.array(allies + enemies, dtype=jnp.uint8)
     return Scenario(
-        place, terrain, unit_starting_sectors, unit_types, n_allies, n_enemies  # type: ignore
+        place,
+        terrain,
+        unit_starting_sectors,  # type: ignore
+        unit_types,
+        n_allies,
+        n_enemies,
+    )
+
+
+def scenario_fn(place, size):
+    # scenario function for Noah, cos the one above is confusing
+    terrain = geo.geography_fn(place, size)
+    num_allies = 10
+    num_enemies = 10
+    unit_types = jnp.array([0] * num_allies + [1] * num_enemies, dtype=jnp.uint8)
+    # start units in default sectors
+    unit_starting_sectors = jnp.array([[0, 0, 1, 1]] * (num_allies + num_enemies))
+    return Scenario(
+        place=place,
+        terrain=terrain,
+        unit_starting_sectors=unit_starting_sectors,
+        unit_types=unit_types,
+        num_allies=num_allies,
+        num_enemies=num_enemies,
     )
 
 
@@ -87,7 +119,7 @@ def spawn_fn(rng: jnp.ndarray, units_spawning_sectors):
     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
+        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)
@@ -101,47 +133,65 @@ def sectors_fn(sectors: jnp.ndarray, invalid_spawn_areas: jnp.ndarray):
     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)
+        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) ))
+        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.building.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])
+            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 if "unit_type_velocities" not in kwargs else kwargs["unit_type_velocities"]
+        self.unit_type_velocities = (
+            jnp.array([3.15, 2.25, 4.13, 3.15, 4.13, 3.15]) / 2.5
+            if "unit_type_velocities" not in kwargs
+            else kwargs["unit_type_velocities"]
+        )
         self.unit_type_attack_blasts = jnp.zeros((3,), dtype=jnp.float32)  # TODO: add
         self.max_steps = 200
         self._push_units_away = lambda state, firmness=1: state  # overwrite push units
-        self.spawning_sectors = sectors_fn(self.unit_starting_sectors, scenario.terrain_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))
-
+        self.spawning_sectors = sectors_fn(
+            self.unit_starting_sectors,
+            scenario.terrain.building + scenario.terrain.water,
+        )
+        self.resolution = (
+            jnp.array(jnp.max(self.unit_type_sight_ranges), dtype=jnp.int32) * 2
+        )
+        self.t = jnp.tile(jnp.linspace(0, 1, self.resolution), (2, 1))
 
-    def reset(self, rng: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]: # type: ignore
+    def reset(self, rng: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]:  # type: ignore
         """Environment-specific reset."""
         unit_positions = spawn_fn(rng, self.spawning_sectors)
         unit_teams = jnp.zeros((self.num_agents,))
@@ -161,24 +211,25 @@ class Environment(SMAX):
             time=0,
             terminal=False,
             unit_weapon_cooldowns=unit_weapon_cooldowns,
-            # terrain=self.terrain_raster,
+            # terrain=self.terrain,
         )
         state = self._push_units_away(state)  # type: ignore  could be slow
         obs = self.get_obs(state)
+        # remove world_state from obs
         world_state = self.get_world_state(state)
-        # obs["world_state"] = jax.lax.stop_gradient(world_state)
+        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)
+        # 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
+        # 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
+    def get_obs_unit_list(self, state: State) -> Dict[str, chex.Array]:  # type: ignore
         """Applies observation function to state."""
 
         def get_features(i, j):
@@ -188,9 +239,7 @@ class Environment(SMAX):
             # 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,
+                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(
@@ -205,8 +254,14 @@ class Environment(SMAX):
                 < 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),
+                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,
             )
@@ -240,20 +295,20 @@ class Environment(SMAX):
             pos
             + firmness * jnp.sum(delta_matrix * overlap_term[:, :, None], axis=1) / 2
         )
-        return jnp.where(self.unit_type_pushable[unit_types][:, None], unit_positions, pos)
-
-    def has_line_of_sight(self, source, target, raster_input):  # 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)
+        return jnp.where(
+            self.unit_type_pushable[unit_types][:, None], unit_positions, pos
+        )
 
+    def has_line_of_sight(self, source, target, raster_input):
+        # suppose that the target is in sight_range of source, otherwise the line of sight might miss some cells
+        cells = jnp.array(
+            source[:, jnp.newaxis] * self.t + (1 - self.t) * target[:, jnp.newaxis],
+            dtype=jnp.int32,
+        )
+        mask = jnp.zeros(raster_input.shape).at[cells[0, :], cells[1, :]].set(1)
         flag = ~jnp.any(jnp.logical_and(mask, raster_input))
-
         return flag
 
-
     @partial(jax.jit, static_argnums=(0,))  # replace the _world_step method
     def _world_step(  # modified version of JaxMARL's SMAX _world_step
         self,
@@ -265,13 +320,12 @@ class Environment(SMAX):
             pos, new_pos = pos.astype(jnp.int32), new_pos.astype(jnp.int32)
             minimum = jnp.minimum(pos, new_pos)
             maximum = jnp.maximum(pos, new_pos)
-            mask = jnp.where(jnp.arange(mask.shape[0]) >= minimum[0], mask, 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
+            mask = jnp.where(jnp.arange(mask.shape[0]) >= minimum[0], mask.T, 0).T
+            mask = jnp.where(jnp.arange(mask.shape[0]) <= maximum[0], mask.T, 0).T
+            mask = jnp.where(jnp.arange(mask.shape[1]) >= minimum[1], mask, 0)
+            mask = jnp.where(jnp.arange(mask.shape[1]) <= maximum[1], mask, 0)
             return jnp.any(mask)
 
-
         def update_position(idx, vec):
             # Compute the movements slightly strangely.
             # The velocities below are for diagonal directions
@@ -287,13 +341,17 @@ class Environment(SMAX):
             )
             # avoid going out of bounds
             new_pos = jnp.maximum(
-                jnp.minimum(new_pos, jnp.array([self.map_width-1, self.map_height-1])),
+                jnp.minimum(
+                    new_pos, jnp.array([self.map_width - 1, self.map_height - 1])
+                ),
                 jnp.zeros((2,)),
             )
 
             #######################################################################
             ############################################ avoid going into obstacles
-            clash = raster_crossing(pos, new_pos, self.terrain_raster.building + self.terrain_raster.water)
+            clash = raster_crossing(
+                pos, new_pos, self.terrain.building + self.terrain.water
+            )
             new_pos = jnp.where(clash, pos, new_pos)
 
             #######################################################################
@@ -331,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]
             )
@@ -349,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]]
+            # )
 
             #########################################################
             #########################################################
@@ -387,26 +449,28 @@ 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, in_axes=(0, 0, None))(pos, new_pos, self.terrain_raster.building + self.terrain_raster.water)
+        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(
@@ -441,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
 
         #########################################################
         #########################################################
@@ -456,14 +520,23 @@ class Environment(SMAX):
         )
         return state
 
+
 if __name__ == "__main__":
     n_envs = 4
 
-
     n_allies = 10
-    scenario_kwargs = {"allies_type": 0, "n_allies": n_allies, "enemies_type": 0, "n_enemies": n_allies,
-                        "place": "Vesterbro, Copenhagen, Denmark", "size": 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_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))
@@ -471,24 +544,18 @@ if __name__ == "__main__":
     obs, state = vmap(env.reset)(reset_key)
     state_seq = []
 
-
-    from tqdm import tqdm
     import time
+
     step = vmap(jit(env.step))
     tic = time.time()
-    for i in tqdm(range(10)):
+    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))
-        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))
-        print(tic - time.time())
         obs, state, reward, done, infos = step(step_key, state, act)
-        print(tic - time.time())
         tic = time.time()

parabellum/geo.py

@@ -21,6 +21,7 @@ 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]
@@ -31,7 +32,20 @@ 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}
+tags = {
+    "building": True,
+    "water": True,
+    "highway": True,
+    "landuse": [
+        "grass",
+        "forest",
+        "flowerbed",
+        "greenfield",
+        "village_green",
+        "recreation_ground",
+    ],
+    "leisure": "garden",
+}  #  "road": True}
 
 
 # %% Coordinate function
@@ -54,34 +68,47 @@ def get_bbox(place: str, buffer) -> BBox:
 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
+    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()
+    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):
+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")
+    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)
+    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 ["building", "water", "landuse"]])
+    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)
@@ -89,12 +116,15 @@ def feature_fn(t, feature, gdf, shape):
     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)
+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/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

@@ -1,3 +1,4 @@
+# %%
 # tps.py
 #   parabellum types and dataclasses
 # by: Noah Syrkis

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: parabellum
-Version: 0.3.0
+Version: 0.3.3
 Summary: Parabellum environment for parallel warfare simulation
 Home-page: https://github.com/syrkis/parabellum
 License: MIT
@@ -38,7 +38,7 @@ 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)
 
@@ -95,3 +95,4 @@ Full documentation: [parabellum.readthedocs.io](https://parabellum.readthedocs.i
 ## License
 
 MIT
+

parabellum-0.3.3.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=yf5pSmTuFEHS4J4zKTprNTVsOni6bABP9IfWTMrO0OU,22581
+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.3.dist-info/METADATA,sha256=OzXtMvFmkyMwAv4d3X7YFRAhDLuiuRda2ytgsgAXDIA,2707
+parabellum-0.3.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+parabellum-0.3.3.dist-info/RECORD,,

parabellum/map.py

@@ -1,95 +0,0 @@
-# ludens.py
-#    script for fucking around and finding out
-# by: Noah Syrkis
-
-
-# %% 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 jax.numpy as jnp
-import geopandas as gpd
-import rasterio
-from rasterio import features
-from shapely.geometry import Point
-from typing import List
-
-# %% Constants
-geolocator = Nominatim(user_agent="parabellum")
-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-0.3.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-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,,