parabellum 0.0.0__py3-none-any.whl → 0.0.73__py3-none-any.whl

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/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.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,8 @@
+parabellum/__init__.py,sha256=yl3tJXQYNnBgAviJINdHrcALb1177n0CFZ2lGiacbQA,109
+parabellum/env.py,sha256=8KaqUwAmP9FXrINOxGWwZmQJeM6eigeWdjNhlWUNi68,3236
+parabellum/geo.py,sha256=lZ4TQvpfPGW3LV9X-1lfWh3LnwoRCtl_DkN6h26xMng,6922
+parabellum/types.py,sha256=mkBwCA-2_syJAZBeNie9iCG9Cy0SvANMAjCo93XIS1A,3981
+parabellum/utils.py,sha256=wOE2nfcdSvGOBFaGlTEIzS6w451arpT9VSOf_3Zhjwc,2497
+parabellum-0.0.73.dist-info/METADATA,sha256=jHqmHqdrCHM2tIam9bMZ7xwvR1BFLy9bc9qSviKcYAc,1481
+parabellum-0.0.73.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+parabellum-0.0.73.dist-info/RECORD,,

{parabellum-0.0.0.dist-info → parabellum-0.0.73.dist-info}/WHEEL

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

parabellum/.ipynb_checkpoints/__init__-checkpoint.py

@@ -1,4 +0,0 @@
-from .env import Parabellum, Scenario
-from .vis import Visualizer
-
-__all__ = ["Parabellum", "Visualizer", "Scenario"]

parabellum/.ipynb_checkpoints/env-checkpoint.py

@@ -1,296 +0,0 @@
-"""Parabellum environment based on SMAX"""
-
-import jax.numpy as jnp
-import jax
-import numpy as np
-from jax import random
-from jax import jit
-from flax.struct import dataclass
-import chex
-from jaxmarl.environments.smax.smax_env import State, SMAX
-from typing import Tuple, Dict
-from functools import partial
-
-
-@dataclass
-class Scenario:
-    """Parabellum scenario"""
-
-    obstacle_coords: chex.Array
-    obstacle_deltas: chex.Array
-
-    unit_types: chex.Array
-    num_allies: int
-    num_enemies: int
-
-    smacv2_position_generation: bool = False
-    smacv2_unit_type_generation: bool = False
-
-
-# default scenario
-scenarios = {
-    "default": Scenario(
-        jnp.array([[6, 10], [26, 10]]) * 8,
-        jnp.array([[0, 12], [0, 1]]) * 8,
-        jnp.zeros((19,), dtype=jnp.uint8),
-        9,
-        10,
-    )
-}
-
-
-class Parabellum(SMAX):
-    def __init__(
-        self,
-        scenario: Scenario = scenarios["default"],
-        unit_type_attack_blasts=jnp.array([0, 0, 0, 0, 0, 0]) + 8,
-        **kwargs,
-    ):
-        super().__init__(scenario=scenario, **kwargs)
-        self.unit_type_attack_blasts = unit_type_attack_blasts
-        self.obstacle_coords = scenario.obstacle_coords.astype(jnp.float32)
-        self.obstacle_deltas = scenario.obstacle_deltas.astype(jnp.float32)
-        self.max_steps = 200
-        # overwrite supers _world_step method
-
-    
-    def _push_units_away(self, state: State, firmness: float = 1.0): # we do it inside the _world_step to allow more obstacles constraints
-        return state
-
-    def _our_push_units_away(self, pos, unit_types, firmness: float = 1.0): # copy of SMAX._push_units_away but used without state and called inside _world_step to allow more obstacles constraints
-        delta_matrix = pos[:, None] - pos[None, :]
-        dist_matrix = (
-            jnp.linalg.norm(delta_matrix, axis=-1)
-            + jnp.identity(self.num_agents)
-            + 1e-6
-        )
-        radius_matrix = (
-            self.unit_type_radiuses[unit_types][:, None]
-            + self.unit_type_radiuses[unit_types][None, :]
-        )
-        overlap_term = jax.nn.relu(radius_matrix / dist_matrix - 1.0)
-        unit_positions = (
-            pos
-            + firmness * jnp.sum(delta_matrix * overlap_term[:, :, None], axis=1) / 2
-        )
-        return unit_positions
-        
-    @partial(jax.jit, static_argnums=(0,))  # replace the _world_step method
-    def _world_step(  # modified version of JaxMARL's SMAX _world_step
-        self,
-        key: chex.PRNGKey,
-        state: State,
-        actions: Tuple[chex.Array, chex.Array],
-    ) -> Tuple[Dict[str, chex.Array], State, Dict[str, float], Dict[str, bool], Dict]:
-        
-        @partial(jax.vmap, in_axes=(None, None, 0, 0))
-        def inter_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 update_position(idx, vec):
-            # Compute the movements slightly strangely.
-            # The velocities below are for diagonal directions
-            # because these are easier to encode as actions than the four
-            # diagonal directions. Then rotate the velocity 45
-            # degrees anticlockwise to compute the movement.
-            pos = state.unit_positions[idx]
-            new_pos = (
-                pos
-                + vec
-                * self.unit_type_velocities[state.unit_types[idx]]
-                * self.time_per_step
-            )
-            # avoid going out of bounds
-            new_pos = jnp.maximum(
-                jnp.minimum(new_pos, jnp.array([self.map_width, self.map_height])),
-                jnp.zeros((2,)),
-            )
-
-            #######################################################################
-            ############################################ avoid going into obstacles
-            obs = self.obstacle_coords
-            obs_end = obs + self.obstacle_deltas
-            inters = jnp.any(inter_fn(pos, new_pos, obs, obs_end))
-            new_pos = jnp.where(inters, pos, new_pos)
-
-            #######################################################################
-            #######################################################################
-
-            return new_pos
-
-        #######################################################################
-        ######################################### units close enough to get hit
-
-        def bystander_fn(attacked_idx):
-            idxs = jnp.zeros((self.num_agents,))
-            idxs *= (
-                jnp.linalg.norm(
-                    state.unit_positions - state.unit_positions[attacked_idx], axis=-1
-                )
-                < self.unit_type_attack_blasts[state.unit_types[attacked_idx]]
-            )
-            return idxs
-
-            #######################################################################
-            #######################################################################
-
-        def update_agent_health(idx, action, key):  # TODO: add attack blasts
-            # for team 1, their attack actions are labelled in
-            # reverse order because that is the order they are
-            # observed in
-            attacked_idx = jax.lax.cond(
-                idx < self.num_allies,
-                lambda: action + self.num_allies - self.num_movement_actions,
-                lambda: self.num_allies - 1 - (action - self.num_movement_actions),
-            )
-            # deal with no-op attack actions (i.e. agents that are moving instead)
-            attacked_idx = jax.lax.select(
-                action < self.num_movement_actions, idx, attacked_idx
-            )
-
-            attack_valid = (
-                (
-                    jnp.linalg.norm(
-                        state.unit_positions[idx] - state.unit_positions[attacked_idx]
-                    )
-                    < self.unit_type_attack_ranges[state.unit_types[idx]]
-                )
-                & state.unit_alive[idx]
-                & state.unit_alive[attacked_idx]
-            )
-            attack_valid = attack_valid & (idx != attacked_idx)
-            attack_valid = attack_valid & (state.unit_weapon_cooldowns[idx] <= 0.0)
-            health_diff = jax.lax.select(
-                attack_valid,
-                -self.unit_type_attacks[state.unit_types[idx]],
-                0.0,
-            )
-            # 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_)
-                .astype(jnp.float32)
-            )
-            bystander_health_diff = (
-                bystander_valid * -self.unit_type_attacks[state.unit_types[idx]]
-            )
-
-            #########################################################
-            #########################################################
-
-            cooldown_deviation = jax.random.uniform(
-                key, minval=-self.time_per_step, maxval=2 * self.time_per_step
-            )
-            cooldown = (
-                self.unit_type_weapon_cooldowns[state.unit_types[idx]]
-                + cooldown_deviation
-            )
-            cooldown_diff = jax.lax.select(
-                attack_valid,
-                # subtract the current cooldown because we are
-                # going to add it back. This way we effectively
-                # set the new cooldown to `cooldown`
-                cooldown - state.unit_weapon_cooldowns[idx],
-                -self.time_per_step,
-            )
-            return (
-                health_diff,
-                attacked_idx,
-                cooldown_diff,
-                (bystander_health_diff, bystander_idxs),
-            )
-
-        def perform_agent_action(idx, action, key):
-            movement_action, attack_action = action
-            new_pos = update_position(idx, movement_action)
-            health_diff, attacked_idxes, cooldown_diff, (bystander) = (
-                update_agent_health(idx, attack_action, key)
-            )
-
-            return new_pos, (health_diff, attacked_idxes), cooldown_diff, bystander
-
-        keys = jax.random.split(key, num=self.num_agents)
-        pos, (health_diff, attacked_idxes), cooldown_diff, bystander = jax.vmap(
-            perform_agent_action
-        )(jnp.arange(self.num_agents), actions, keys)
-
-        # checked that no unit passed through an obstacles 
-        new_pos = self._our_push_units_away(pos, state.unit_types)
-
-        # avoid going into obstacles after being pushed 
-        obs = self.obstacle_coords
-        obs_end = obs + self.obstacle_deltas
-   
-        def check_obstacles(pos, new_pos, obs, obs_end):
-            inters = jnp.any(inter_fn(pos, new_pos, obs, obs_end))
-            return jnp.where(inters, pos, new_pos)
-        
-        pos = jax.vmap(check_obstacles, in_axes=(0,0,None,None))(pos, new_pos, obs, obs_end)
-        
-        # Multiple enemies can attack the same unit.
-        # We have `(health_diff, attacked_idx)` pairs.
-        # `jax.lax.scatter_add` aggregates these exactly
-        # in the way we want -- duplicate idxes will have their
-        # health differences added together. However, it is a
-        # super thin wrapper around the XLA scatter operation,
-        # which has this bonkers syntax and requires this dnums
-        # parameter. The usage here was inferred from a test:
-        # https://github.com/google/jax/blob/main/tests/lax_test.py#L2296
-        dnums = jax.lax.ScatterDimensionNumbers(
-            update_window_dims=(),
-            inserted_window_dims=(0,),
-            scatter_dims_to_operand_dims=(0,),
-        )
-        unit_health = jnp.maximum(
-            jax.lax.scatter_add(
-                state.unit_health,
-                jnp.expand_dims(attacked_idxes, 1),
-                health_diff,
-                dnums,
-            ),
-            0.0,
-        )
-
-        #########################################################
-        ############################ subtracting bystander health
-
-        _, bystander_health_diff = bystander
-        unit_health -= bystander_health_diff.sum(axis=0)  # might be axis=1
-
-        #########################################################
-        #########################################################
-
-        unit_weapon_cooldowns = state.unit_weapon_cooldowns + cooldown_diff
-        state = state.replace(
-            unit_health=unit_health,
-            unit_positions=pos,
-            unit_weapon_cooldowns=unit_weapon_cooldowns,
-        )
-        return state
-
-if __name__ == "__main__":
-    env = Parabellum(map_width=256, map_height=256)
-    rng, key = random.split(random.PRNGKey(0))
-    obs, state = env.reset(key)
-    state_seq = []
-    for step in range(100):
-        rng, key = random.split(rng)
-        key_act = random.split(key, len(env.agents))
-        actions = {
-            agent: jax.random.randint(key_act[i], (), 0, 5)
-            for i, agent in enumerate(env.agents)
-        }
-        _, state, _, _, _ = env.step(key, state, actions)
-        state_seq.append((obs, state, actions))
-
-