parabellum 0.0.0__tar.gz → 0.0.73__tar.gz

parabellum-0.0.73/.gitignore

@@ -0,0 +1,16 @@
+**/__pycache__
+**/*.gif
+**/*.svg
+*.mp4
+dist/
+**/.DS_Store
+**/.ipynb_checkpoints
+cache/
+pyrightconfig.json
+.venv
+output
+parabellum/cache/
+.ruff_cache/
+.coverage
+build/
+profiler_logs/

parabellum-0.0.73/PKG-INFO

@@ -0,0 +1,46 @@
+Metadata-Version: 2.4
+Name: parabellum
+Version: 0.0.73
+Summary: Parabellum environment for parallel warfare simulation
+Author-email: Noah Syrkis <desk@syrkis.com>
+Requires-Python: <3.12,>=3.11
+Requires-Dist: brax<0.13,>=0.12.1
+Requires-Dist: cachier<4,>=3.1.2
+Requires-Dist: cartopy<0.24,>=0.23.0
+Requires-Dist: contextily<2,>=1.6.0
+Requires-Dist: distrax<0.2,>=0.1.5
+Requires-Dist: einops<0.9,>=0.8.0
+Requires-Dist: equinox>=0.12.2
+Requires-Dist: evosax<0.2,>=0.1.6
+Requires-Dist: flashbax<0.2,>=0.1.2
+Requires-Dist: flax<0.11,>=0.10.4
+Requires-Dist: folium<0.18,>=0.17.0
+Requires-Dist: geopy<3,>=2.4.1
+Requires-Dist: gymnax<0.0.9,>=0.0.8
+Requires-Dist: ipykernel<7,>=6.29.5
+Requires-Dist: ipython>=8.36.0
+Requires-Dist: jax-tqdm<0.4,>=0.3.1
+Requires-Dist: jax<0.7,>=0.6.0
+Requires-Dist: jaxkd>=0.1.0
+Requires-Dist: jaxtyping<0.3,>=0.2.33
+Requires-Dist: jupyterlab<5,>=4.2.2
+Requires-Dist: navix<0.8,>=0.7.0
+Requires-Dist: notebook>=7.4.2
+Requires-Dist: numpy>=2
+Requires-Dist: omegaconf<3,>=2.3.0
+Requires-Dist: optax<0.3,>=0.2.4
+Requires-Dist: osmnx==2.0.0b0
+Requires-Dist: pandas<3,>=2.2.2
+Requires-Dist: poetry<2,>=1.8.3
+Requires-Dist: rasterio<2,>=1.3.10
+Requires-Dist: stadiamaps<4,>=3.2.1
+Requires-Dist: tensorboard>=2.19.0
+Requires-Dist: tensorflow>=2.19.0
+Requires-Dist: tqdm<5,>=4.66.4
+Requires-Dist: wandb<0.20,>=0.19.7
+Requires-Dist: xprof>=2.20.0
+Description-Content-Type: text/markdown
+
+# Parabellum
+
+TODO: switch to red and blue team semantics (not enemy and ally)

parabellum-0.0.73/README.md

@@ -0,0 +1,3 @@
+# Parabellum
+
+TODO: switch to red and blue team semantics (not enemy and ally)

parabellum-0.0.73/main.py

@@ -0,0 +1,57 @@
+# %% main.py
+#    parabellum main
+# by: Noah Syrkis
+
+# Imports ###################################################################
+from functools import partial
+from typing import Tuple
+import esch
+
+from einops import repeat, rearrange
+import numpy as np
+from jax import jit, lax, random, tree, vmap
+from jax import numpy as jnp
+from jaxtyping import Array
+
+import parabellum as pb
+from parabellum.types import Action, Config, State
+
+
+# %% Functions
+def action_fn(cfg: Config, rng: Array) -> Action:
+    pos = random.uniform(rng, (cfg.length, 2), minval=-1, maxval=1) * cfg.rules.reach[cfg.types][..., None]
+    kind = random.randint(rng, (cfg.length,), minval=0, maxval=3)
+    return Action(pos=pos, kind=kind)
+
+
+def step_fn(state: State, rng: Array) -> Tuple[State, Tuple[State, Action]]:
+    action = action_fn(cfg, rng)
+    obs, state = env.step(cfg, rng, state, action)
+    return state, (state, action)
+
+
+def traj_fn(state, rng) -> Tuple[State, Tuple[State, Action]]:
+    rngs = random.split(rng, cfg.steps)
+    return lax.scan(step_fn, state, rngs)
+
+
+# %% Main
+env, cfg = pb.env.Env(), Config()
+init_key, traj_key = random.split(random.PRNGKey(0), (2, cfg.sims))
+
+init = vmap(jit(partial(env.init, cfg)))
+traj = vmap(jit(traj_fn))
+
+obs, state = init(init_key)
+state, (seq, action) = traj(state, init_key)
+
+pb.utils.svg_fn(cfg, seq, action, "/Users/nobr/desk/s3/parabellum/sims.svg", debug=True)
+# %% Anim
+# for i in range(seq.pos.shape[0]):  # sims
+# for j in range(seq.pos.shape[2]):  # units
+# shots = [(kdx, coord) for kdx, coord in enumerate(action.pos[i, :, j]) if action.shoot[i, kdx, j]]
+# print(shots)
+# print(tree.map(jnp.shape, action))
+# print(tree.map(jnp.shape, seq))
+# pb.utils.svg_fn(cfg, seq.pos, action)
+# pb.utils.gif_fn(cfg, seq)

parabellum-0.0.73/parabellum/__init__.py

@@ -0,0 +1,4 @@
+from . import env, geo, types, utils
+from .env import Env
+
+__all__ = ["geo", "env", "types", "utils", "Env"]

parabellum-0.0.73/parabellum/env.py

@@ -0,0 +1,70 @@
+# %% env.py
+#   parabellum env
+# by: Noah Syrkis
+
+# Imports
+from functools import partial
+from typing import Tuple
+
+import jax.numpy as jnp
+import jaxkd as jk
+from jax import random
+from jaxtyping import Array
+
+from parabellum.types import Action, Config, Obs, State
+
+
+# %% Dataclass
+class Env:
+    def init(self, cfg: Config, rng: Array) -> Tuple[Obs, State]:
+        state = init_fn(cfg, rng)  # without jit this takes forever
+        return obs_fn(cfg, state), state
+
+    def step(self, cfg: Config, rng: Array, state: State, action: Action) -> Tuple[Obs, State]:
+        state = step_fn(cfg, rng, state, action)
+        return obs_fn(cfg, state), state
+
+
+# %% Functions
+def init_fn(cfg: Config, rng: Array) -> State:
+    prob = jnp.ones((cfg.size, cfg.size)).at[cfg.map].set(0).flatten()  # Set
+    flat = random.choice(rng, jnp.arange(prob.size), shape=(cfg.length,), p=prob, replace=True)
+    idxs = (flat // len(cfg.map), flat % len(cfg.map))
+    pos = jnp.float32(jnp.column_stack(idxs))
+    return State(pos=pos, hp=cfg.hp[cfg.types])
+
+
+def obs_fn(cfg: Config, state: State) -> Obs:  # return info about neighbors ---
+    idxs, dist = jk.extras.query_neighbors_pairwise(state.pos, state.pos, k=cfg.knn)
+    mask = dist < cfg.sight[cfg.types[idxs][:, 0]][..., None]  # | (state.hp[idxs] > 0)
+    pos = (state.pos[idxs] - state.pos[:, None, ...]).at[:, 0, :].set(state.pos) * mask[..., None]
+    args = state.hp, cfg.types, cfg.teams, cfg.reach, cfg.sight, cfg.speed
+    hp, type, team, reach, sight, speed = map(lambda x: x[idxs] * mask, args)
+    return Obs(pos=pos, dist=dist, hp=hp, type=type, team=team, reach=reach, sight=sight, speed=speed, mask=mask)
+
+
+def step_fn(cfg: Config, rng: Array, state: State, action: Action) -> State:
+    idx, norm = jk.extras.query_neighbors_pairwise(state.pos + action.pos, state.pos, k=2)
+    args = rng, cfg, state, action, idx, norm
+    return State(pos=partial(push_fn, cfg, rng, idx, norm)(move_fn(*args)), hp=blast_fn(*args))  # type: ignore
+
+
+def move_fn(rng: Array, cfg: Config, state: State, action: Action, idx: Array, norm: Array) -> Array:
+    speed = cfg.speed[cfg.types][..., None]  # max speed of a unit (step size, really)
+    pos = state.pos + action.pos.clip(-speed, speed) * action.move[..., None]  # new poss
+    mask = ((pos < 0).any(axis=-1) | ((pos >= cfg.size).any(axis=-1)) | (cfg.map[*jnp.int32(pos).T] > 0))[..., None]
+    return jnp.where(mask, state.pos, pos)  # compute new position
+
+
+def blast_fn(rng: Array, cfg: Config, state: State, action: Action, idx: Array, norm: Array) -> Array:
+    dam = (cfg.dam[cfg.types] * action.cast)[..., None] * jnp.ones_like(idx)
+    return state.hp - jnp.zeros(cfg.length, dtype=jnp.int32).at[idx.flatten()].add(dam.flatten())
+
+
+def push_fn(cfg: Config, rng: Array, idx: Array, norm: Array, pos: Array) -> Array:
+    return pos + random.normal(rng, pos.shape) * 0.1
+    # params need to be tweaked, and matched with unit size
+    pos_diff = pos[:, None, :] - pos[idx]  # direction away from neighbors
+    mask = (norm < cfg.r[cfg.types][..., None]) & (norm > 0)
+    pos = pos + jnp.where(mask[..., None], pos_diff * cfg.force / (norm[..., None] + 1e-6), 0.0).sum(axis=1)
+    return pos + random.normal(rng, pos.shape) * 0.1

parabellum-0.0.73/parabellum/geo.py

@@ -0,0 +1,213 @@
+# %% geo.py
+#   script for geospatial level generation
+# by: Noah Syrkis
+
+# %% Imports
+from rasterio import features, transform
+
+# from jax import tree
+from geopy.geocoders import Nominatim
+from geopy.distance import distance
+import contextily as cx
+import jax.numpy as jnp
+import cartopy.crs as ccrs
+from jaxtyping import Array
+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
+from cachier import cachier
+# from jax.scipy.signal import convolve
+# from parabellum.types import Terrain
+
+# %% 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)  # type: ignore
+
+
+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
+
+
+@cachier()
+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 = jnp.rot90(basemap_fn(bbox, gdf), 3)
+    # kernel = jnp.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
+    trans = lambda x: jnp.bool(x)  # jnp.rot90(x, 3)  # noqa
+    terrain = trans(raster[0])  # 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,
+    #    )
+    # terrain = tree.map(lambda x: x.astype(jnp.int16), terrain)
+    return terrain
+
+
+# =======
+#     terrain = tps.Terrain(building=trans(raster[0] - convolve(raster[0]*raster[2], kernel, mode='same')>0),
+#                           water=trans(raster[1] - convolve(raster[1]*raster[2], kernel, mode='same')>0),
+#                           forest=trans(jnp.logical_or(raster[3], raster[4])),
+#                           basemap=basemap)
+#     return terrain, gdf
+# >>>>>>> aeb13033e57083cc512a60f8f60a3db47a65ac32
+
+
+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
+
+
+# %%
+# def normalize(x):
+#     return (np.array(x) - m) / (M - m)
+
+
+# def get_bridges(gdf):
+#     xmin, ymin, xmax, ymax = gdf.total_bounds
+#     m = np.array([xmin, ymin])
+#     M = np.array([xmax, ymax])
+
+#     bridges = {}
+#     for idx, bridge in gdf[gdf["bridge"] == "yes"].iterrows():
+#         if type(bridge["name"]) == str:
+#             bridges[idx[1]] = {
+#                 "name": bridge["name"],
+#                 "coords": normalize(
+#                     [bridge.geometry.centroid.x, bridge.geometry.centroid.y]
+#                 ),
+#             }
+#     return bridges
+
+
+"""
+# %%
+if __name__ == "__main__":
+    place = "Thun, Switzerland"
+<<<<<<< HEAD
+    terrain = geography_fn(place, 300)
+
+=======
+    terrain, gdf = geography_fn(place, 300)
+
+>>>>>>> aeb13033e57083cc512a60f8f60a3db47a65ac32
+    fig, axes = plt.subplots(1, 5, figsize=(20, 20))
+    axes[0].imshow(jnp.rot90(terrain.building), cmap="gray")
+    axes[1].imshow(jnp.rot90(terrain.water), cmap="gray")
+    axes[2].imshow(jnp.rot90(terrain.forest), cmap="gray")
+    axes[3].imshow(jnp.rot90(terrain.building + terrain.water + terrain.forest))
+    axes[4].imshow(jnp.rot90(terrain.basemap))
+
+    # %%
+    W, H, _  = terrain.basemap.shape
+    bridges = get_bridges(gdf)
+
+    # %%
+    print("Bridges:")
+    for bridge in bridges.values():
+        x, y = int(bridge["coords"][0]*300), int(bridge["coords"][1]*300)
+        print(bridge["name"], f"at ({x}, {y})")
+
+    # %%
+    plt.subplots(figsize=(7,7))
+    plt.imshow(jnp.rot90(terrain.basemap))
+    X = [b["coords"][0]*W for b in bridges.values()]
+    Y = [(1-b["coords"][1])*H for b in bridges.values()]
+    plt.scatter(X, Y)
+    for i in range(len(X)):
+        x,y = int(X[i]), int(Y[i])
+        plt.text(x, y, str((int(x/W*300), int((1-(y/H))*300))))
+
+# %%
+
+# %% [raw]
+# fig, ax = plt.subplots(figsize=(10, 10))
+# gdf.plot(ax=ax, color='lightgray')  # Plot all features
+# bridges.plot(ax=ax, color='red')     # Highlight bridges in red
+# plt.show()
+
+# %%
+
+"""
+
+# BBox = namedtuple("BBox", ["north", "south", "east", "west"])  # type: ignore
+
+
+# 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-0.0.73/parabellum/types.py

@@ -0,0 +1,138 @@
+# types.py
+#   parabellum types
+# by: Noah Syrkis
+
+# imports
+from chex import dataclass
+from jaxtyping import Array, Float32, Int
+import jax.numpy as jnp
+from parabellum.geo import geography_fn
+from dataclasses import field
+
+
+@dataclass
+class Kind:
+    hp: int
+    dam: int
+    speed: int
+    reach: int
+    sight: int
+    blast: int
+    r: float
+
+
+@dataclass
+class Rules:
+    troop = Kind(hp=120, dam=15, speed=2, reach=4, sight=4, blast=1, r=1)
+    armor = Kind(hp=150, dam=12, speed=1, reach=8, sight=16, blast=3, r=2)
+    plane = Kind(hp=80, dam=20, speed=4, reach=16, sight=32, blast=2, r=2)
+    civil = Kind(hp=100, dam=0, speed=3, reach=5, sight=10, blast=1, r=2)
+    medic = Kind(hp=100, dam=-10, speed=3, reach=5, sight=10, blast=1, r=2)
+
+    def __post_init__(self):
+        self.hp = jnp.array((self.troop.hp, self.armor.hp, self.plane.hp, self.civil.hp, self.medic.hp))
+        self.dam = jnp.array((self.troop.dam, self.armor.dam, self.plane.dam, self.civil.dam, self.medic.dam))
+        self.r = jnp.array((self.troop.r, self.armor.r, self.plane.r, self.civil.r, self.medic.r))
+        self.speed = jnp.array(
+            (self.troop.speed, self.armor.speed, self.plane.speed, self.civil.speed, self.medic.speed)
+        )
+        self.reach = jnp.array(
+            (self.troop.reach, self.armor.reach, self.plane.reach, self.civil.reach, self.medic.reach)
+        )
+        self.sight = jnp.array(
+            (self.troop.sight, self.armor.sight, self.plane.sight, self.civil.sight, self.medic.sight)
+        )
+        self.blast = jnp.array(
+            (self.troop.blast, self.armor.blast, self.plane.blast, self.civil.blast, self.medic.blast)
+        )
+
+
+@dataclass
+class Team:
+    troop: int = 1
+    armor: int = 0
+    plane: int = 0
+    civil: int = 0
+    medic: int = 0
+
+    def __post_init__(self):
+        self.length: int = self.troop + self.armor + self.plane + self.civil + self.medic
+        self.types: Array = jnp.repeat(
+            jnp.arange(5), jnp.array((self.troop, self.armor, self.plane, self.civil, self.medic))
+        )
+
+
+# dataclasses
+@dataclass
+class State:
+    pos: Array
+    hp: Array
+    # target: Array
+
+
+@dataclass
+class Obs:
+    # idxs: Array
+    hp: Array
+    pos: Array
+    type: Array
+    team: Array
+    dist: Array
+    mask: Array
+    reach: Array
+    sight: Array
+    speed: Array
+
+    @property
+    def ally(self):
+        return (self.team == self.team[0]) & self.mask
+
+    @property
+    def enemy(self):
+        return (self.team != self.team[0]) & self.mask
+
+
+@dataclass
+class Action:
+    pos: Array
+    kind: Int[Array, "..."]  # 0 = invalid, 1 = move, 2 = cast
+
+    @property
+    def invalid(self):
+        return self.kind == 0
+
+    @property
+    def move(self):
+        return self.kind == 1
+
+    @property
+    def cast(self):  # cast bomb, bullet or medicin
+        return self.kind == 2
+
+
+@dataclass
+class Config:  # Remove frozen=True for now
+    steps: int = 123
+    place: str = "Palazzo della Civiltà Italiana, Rome, Italy"
+    force: float = 0.5
+    sims: int = 2
+    size: int = 64
+    knn: int = 2
+    blu: Team = field(default_factory=lambda: Team())
+    red: Team = field(default_factory=lambda: Team())
+    rules: Rules = field(default_factory=lambda: Rules())
+
+    def __post_init__(self):
+        # Pre-compute everything once
+        self.types: Array = jnp.concat((self.blu.types, self.red.types))
+        self.teams: Array = jnp.repeat(jnp.arange(2), jnp.array((self.blu.length, self.red.length)))
+        self.map: Array = geography_fn(self.place, self.size)  # Computed once here
+        self.hp: Array = self.rules.hp
+        self.dam: Array = self.rules.dam
+        self.r: Array = self.rules.r
+        self.speed: Array = self.rules.speed
+        self.reach: Array = self.rules.reach
+        self.sight: Array = self.rules.sight
+        self.blast: Array = self.rules.blast
+        self.length: int = self.blu.length + self.red.length
+        self.root: Array = jnp.int32(jnp.sqrt(self.length))

parabellum-0.0.73/parabellum/utils.py

@@ -0,0 +1,59 @@
+# %% utils.py
+#   parabellum ut
+
+
+# Imports
+import esch
+import jax.numpy as jnp
+import numpy as np
+from einops import rearrange, repeat
+from jax import tree
+from PIL import Image
+from parabellum.types import Config
+
+# Twilight colors (used in neurocope)
+red = "#EA344A"
+blue = "#2B60F6"
+
+
+# %% Plotting
+def gif_fn(cfg: Config, seq, scale=4):  # animate positions TODO: remove dead units
+    pos = seq.pos.astype(int)
+    cord = jnp.concat((jnp.arange(pos.shape[0]).repeat(pos.shape[1])[..., None], pos.reshape(-1, 2)), axis=1).T
+    idxs = cord[:, seq.hp.flatten().astype(bool) > 0]
+    imgs = 1 - np.array(repeat(cfg.map, "... -> a ...", a=len(pos)).at[*idxs].set(1))
+    imgs = [Image.fromarray(img).resize(np.array(img.shape[:2]) * scale, Image.NEAREST) for img in imgs * 255]  # type: ignore
+    imgs[0].save("/Users/nobr/desk/s3/btc2sim/sims.gif", save_all=True, append_images=imgs[1:], duration=10, loop=0)
+
+
+def svg_fn(cfg: Config, seq, action, fname, targets=None, fps=2, debug=False):
+    # set up and background
+    e = esch.Drawing(h=cfg.size, w=cfg.size, row=1, col=seq.pos.shape[0], debug=debug, pad=10)
+    esch.grid_fn(repeat(np.array(cfg.map, dtype=float), f"... -> {seq.pos.shape[0]} ...") * 0.5, e, shape="square")
+
+    # loop thorugh teams
+    for i in jnp.unique(cfg.teams):  # c#fg.teams.unique():
+        col = "red" if i == 1 else "blue"
+
+        # loop through types
+        for j in jnp.unique(cfg.types):
+            mask = (cfg.teams == i) & (cfg.types == j)
+            size, blast = float(cfg.rules.r[j]), float(cfg.rules.blast[j])
+            subset = np.array(rearrange(seq.pos, "a b c d -> a c d b"), dtype=float)[:, mask]
+            # print(tree.map(jnp.shape, action), mask.shape)
+            sub_action = tree.map(lambda x: x[:, :, mask], action)
+            # print(tree.map(jnp.shape, sub_action))
+            esch.sims_fn(e, subset, action=sub_action, fps=fps, col=col, stroke=col, size=size, blast=blast)
+
+            if debug:
+                sight, reach = float(cfg.rules.sight[j]), float(cfg.rules.reach[j])
+                esch.sims_fn(e, subset, action=None, col="none", fps=fps, size=reach, stroke="grey")
+                esch.sims_fn(e, subset, action=None, col="none", fps=fps, size=sight, stroke="yellow")
+
+    if targets is not None:
+        pos = np.array(repeat(targets, f"... -> {seq.pos.shape[0]} ..."))
+        arr = np.ones(pos.shape[:-1])
+        esch.mesh_fn(e, pos, arr, shape="square", col="purple")
+
+    # save
+    e.dwg.saveas(fname)

parabellum-0.0.73/pyproject.toml

@@ -0,0 +1,63 @@
+[project]
+name = "parabellum"
+version = "0.0.73"
+description = "Parabellum environment for parallel warfare simulation"
+authors = [{ name = "Noah Syrkis", email = "desk@syrkis.com" }]
+requires-python = ">=3.11,<3.12"
+readme = "README.md"
+dependencies = [
+    "jupyterlab>=4.2.2,<5",
+    "poetry>=1.8.3,<2",
+    "tqdm>=4.66.4,<5",
+    "geopy>=2.4.1,<3",
+    "osmnx==2.0.0b0",
+    "rasterio>=1.3.10,<2",
+    "ipykernel>=6.29.5,<7",
+    "folium>=0.17.0,<0.18",
+    "pandas>=2.2.2,<3",
+    "contextily>=1.6.0,<2",
+    "einops>=0.8.0,<0.9",
+    "jaxtyping>=0.2.33,<0.3",
+    "cartopy>=0.23.0,<0.24",
+    "stadiamaps>=3.2.1,<4",
+    "cachier>=3.1.2,<4",
+    "jax>=0.6.0,<0.7",
+    "gymnax>=0.0.8,<0.0.9",
+    "evosax>=0.1.6,<0.2",
+    "distrax>=0.1.5,<0.2",
+    "optax>=0.2.4,<0.3",
+    "flax>=0.10.4,<0.11",
+    "brax>=0.12.1,<0.13",
+    "wandb>=0.19.7,<0.20",
+    "flashbax>=0.1.2,<0.2",
+    "navix>=0.7.0,<0.8",
+    "omegaconf>=2.3.0,<3",
+    "jax-tqdm>=0.3.1,<0.4",
+    "ipython>=8.36.0",
+    "notebook>=7.4.2",
+    "equinox>=0.12.2",
+    "tensorboard>=2.19.0",
+    "tensorflow>=2.19.0",
+    "xprof>=2.20.0",
+    "jaxkd>=0.1.0",
+    "numpy>=2",
+]
+
+[dependency-groups]
+dev = ["esch"]
+
+[tool.uv]
+
+[tool.uv.sources]
+esch = { path = "../../esch" }
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.pyright]
+venvPath = "."
+venv = ".venv"
+
+[tool.ruff]
+line-length = 120