{"id":24539141,"url":"https://github.com/i-m-iron-man/foragax","last_synced_at":"2025-07-27T20:34:37.922Z","repository":{"id":256550506,"uuid":"715512107","full_name":"i-m-iron-man/Foragax","owner":"i-m-iron-man","description":"An Agent-Based Modelling framework in JAX","archived":false,"fork":false,"pushed_at":"2025-01-10T13:35:08.000Z","size":18873,"stargazers_count":4,"open_issues_count":0,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-04-15T04:38:26.123Z","etag":null,"topics":["agent-based","agent-based-modeling","agent-based-simulation","foraging","jax"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/i-m-iron-man.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE.md","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2023-11-07T09:45:00.000Z","updated_at":"2025-04-11T13:27:56.000Z","dependencies_parsed_at":"2024-10-25T22:56:23.086Z","dependency_job_id":"24af2b39-e93f-40ce-a176-616cc5cdcf10","html_url":"https://github.com/i-m-iron-man/Foragax","commit_stats":null,"previous_names":["i-m-iron-man/foragax"],"tags_count":1,"template":false,"template_full_name":null,"purl":"pkg:github/i-m-iron-man/Foragax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-m-iron-man%2FForagax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-m-iron-man%2FForagax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-m-iron-man%2FForagax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-m-iron-man%2FForagax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/i-m-iron-man","download_url":"https://codeload.github.com/i-m-iron-man/Foragax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-m-iron-man%2FForagax/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":267419226,"owners_count":24084126,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-07-27T02:00:11.917Z","response_time":82,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["agent-based","agent-based-modeling","agent-based-simulation","foraging","jax"],"created_at":"2025-01-22T16:15:34.354Z","updated_at":"2025-07-27T20:34:37.901Z","avatar_url":"https://github.com/i-m-iron-man.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cdiv align=\"center\"\u003e\n    \u003cimg src=\"https://github.com/i-m-iron-man/Foragax/blob/main/docs/assets/small_foragaing.gif\" width=\"250\"/\u003e\n    \u003cimg src=\"https://github.com/i-m-iron-man/Foragax/blob/main/docs/assets/foragax_logo.webp\" width=\"250\"/\u003e\n    \u003cimg src=\"https://github.com/i-m-iron-man/Foragax/blob/main/docs/assets/sheep_wolf.gif\" width=\"250\"/\u003e\n\u003c/div\u003e\n\nA more general Agent-based modeling framework is now available at [abmax](https://github.com/i-m-iron-man/abmax). Please consider using that. All further developments will occur there.\n\nForagax is an Agent-Based Modelling (ABM) package based on JAX. It provides scalable and efficient ABM simulations by leveraging JAX's automatic vectorization and just-in-time compilation capabilities. The main features of Foragax include:\n\n - Agent manipulation (adding, removing, updating, selecting, and sorting agents) with just-in-time compilation.\n - Vectorized ray-casting and wall-detection for simulating agents moving in a continuous 2D environment with custom boundaries and obstacles.\n - Tutorials and examples to help users get started with ABM using JAX.\n - Familiar ABM interface for creating and manipulating agents.\n\n## Installation\n```\npip install foragax\n```\nRequires Python 3.10+, [JAX 0.4.13+](https://jax.readthedocs.io/en/latest/quickstart.html), and [flax 0.7.4+](https://flax.readthedocs.io/en/latest/quick_start.html)\n\n## Hello World : Rolling Dice\nThis is an example of agents that roll a dice.\nFor each agent that draws a 6, a new agent is activated.\nEach agent that draws a 1 is deactivated.\nThe note-book implementing this can be found [here](https://github.com/i-m-iron-man/Foragax/blob/main/examples/basic/hello_world/hello_world.ipynb).\n\n```python\nimport foragax.base.agent_classes as fgx_classes\nimport foragax.base.agent_methods as fgx_methods\nimport jax\nimport jax.numpy as jnp\nfrom flax import struct\n\n@struct.dataclass\nclass Dice(fgx_classes.Agent):\n    @staticmethod\n    def create_agent(params: fgx_classes.Params, unique_id: int, active_state: int, agent_type: int, key:jax.random.PRNGKey):\n        key, subkey = jax.random.split(key)\n        \n        def create_active_agent(key):\n            draw = jax.random.randint(key, (1,), 1, 7)\n            state_content = {'draw': draw, 'key': key}\n            return fgx_classes.State(content=state_content)\n        \n        def create_inactive_agent(key):\n            state_content = {'draw': jnp.array([0]), 'key': key}\n            return fgx_classes.State(content=state_content)\n        agent_state = jax.lax.cond(active_state, lambda _: create_active_agent(subkey), lambda _: create_inactive_agent(subkey), None)\n        \n        return Dice(params = params, unique_id = unique_id, agent_type = agent_type, \n                    active_state = active_state, state = agent_state, policy = None, age = 0.0)\n    \n    @staticmethod\n    def step_agent(params: fgx_classes.Params, input: fgx_classes.Signal, dice_agent: fgx_classes.Agent):\n        \n        def step_active_agent(dice_agent):\n            old_state = dice_agent.state.content\n            key, subkey = jax.random.split(old_state['key'])\n            \n            draw = jax.random.randint(subkey, (1,), 1, 7)\n            state_content = {'draw': draw, 'key': subkey}\n            new_state = fgx_classes.State(content = state_content)\n            return dice_agent.replace(state = new_state, age = dice_agent.age + 1.0)\n        \n        def step_inactive_agent(dice_agent):\n            return dice_agent\n        \n        new_dice_agent = jax.lax.cond(dice_agent.active_state, lambda _: step_active_agent(dice_agent), lambda _: step_inactive_agent(dice_agent), None)\n        return new_dice_agent\n    \n    @staticmethod\n    def add_agent(params: fgx_classes.Params, dice_agents: fgx_classes.Agent, idx, key: jax.random.PRNGKey):\n        inactive_dice_agent = jax.tree_util.tree_map(lambda x:x[idx], dice_agents)\n        useless_key, subkey = jax.random.split(inactive_dice_agent.state.content['key'])\n        draw = jax.random.randint(subkey, (1,), 1, 7)\n        state_content = {'draw': draw, 'key': subkey}\n        new_state = fgx_classes.State(content=state_content)\n        active_dice_agent = inactive_dice_agent.replace(active_state = True, state = new_state)\n        return active_dice_agent, key\n    \n    @staticmethod\n    def remove_agent(params: fgx_classes.Params, dice_agents:fgx_classes.Agent, idx):\n        active_dice_agent = jax.tree_util.tree_map(lambda x:x[idx], dice_agents)\n        draw = jnp.array([0])\n        state_content = {'draw': draw, 'key': active_dice_agent.state.content['key']}\n        state = fgx_classes.State(content=state_content)\n        inactive_dice_agent = active_dice_agent.replace(active_state = False, state = state)\n        return inactive_dice_agent\n\n\n\nDice_set = fgx_classes.Agent_Set(agent = Dice, num_total_agents = 10, num_active_agents = 5, agent_type = 0)\n\nDice_set.agents = fgx_methods.create_agents(params = None, agent_set = Dice_set, key = jax.random.PRNGKey(0))\n\nDice_set.agents = fgx_methods.step_agents(params = None, agent_set = Dice_set, input=None)\n\n# remove all agents who have drawn a 1\n# first, select all agents who have drawn a 1\ndef is_one(dice_agent: fgx_classes.Agent, select_params: fgx_classes.Params):\n    draw = jnp.reshape(dice_agent.state.content['draw'], (-1))\n    return draw == 1\nnum_agents_dead, remove_indices = fgx_methods.jit_select_agents(select_func = is_one, select_params = None, agents = Dice_set.agents)\n\n# now, remove the agents\ndice_remove_params_content = {'remove_ids': remove_indices}\ndice_remove_params = fgx_classes.Params(content = dice_remove_params_content)\nDice_set.agents = fgx_methods.jit_remove_agents(remove_func = Dice.remove_agent, num_agents_remove = num_agents_dead, \n                                                remove_params = dice_remove_params, agents = Dice_set.agents)\n\n# sort agents by active state, as new agents are ALWAYS added at the END of the set\nDice_set.agents, sorted_indices = fgx_methods.jit_sort_agents(quantity = Dice_set.agents.active_state, ascend = False, agents = Dice_set.agents)\n\n\n# add a new agent for every agent that has drawn a 6\n# first, select all agents who have drawn a 6\ndef is_six(dice_agent: fgx_classes.Agent, select_params: fgx_classes.Params):\n    draw = jnp.reshape(dice_agent.state.content['draw'], (-1))\n    return draw == 6\nnum_agents_add, add_indices = fgx_methods.jit_select_agents(select_func = is_six, select_params = None, agents = Dice_set.agents)\n\n# clip the number of agents to add to the number of inactive agents\nnum_active_agents = jnp.sum(Dice_set.agents.active_state, dtype = jnp.int32)\nnum_agents_add = jnp.minimum(num_agents_add, Dice_set.num_total_agents - num_active_agents)\n\n# now, add the agents\nDice_set.agents, key = fgx_methods.jit_add_agents(add_func = Dice.add_agent, num_agents_add = num_agents_add, \n                                                  add_params = None, agents = Dice_set.agents, key = None)\n\n\n```\n## Version v0.0.5\nThe framework is still under active development. Feel free to open an issue if you find any bugs or have any suggestions.\n\n## Citation\n\nIf this framework was useful in your work, please consider starring and cite: [(arXiv link)](https://www.arxiv.org/abs/2409.06345v2)\n\n```bibtex\n@misc{chaturvedi2024foragaxagentbasedmodelling,\n      title={Foragax: An Agent Based Modelling framework based on JAX}, \n      author={Siddharth Chaturvedi and Ahmed El-Gazzar and Marcel van Gerven},\n      year={2024},\n      eprint={2409.06345},\n      archivePrefix={arXiv},\n      primaryClass={cs.MA},\n      url={https://arxiv.org/abs/2409.06345}, \n}\n```\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fi-m-iron-man%2Fforagax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fi-m-iron-man%2Fforagax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fi-m-iron-man%2Fforagax/lists"}