# Kozax: Genetic Programming in JAX
Kozax introduces a general framework for evolving computer programs with genetic programming in JAX. With JAX, the computer programs can be vectorized and evaluated on parallel on CPU and GPU. Furthermore, Just-in-time compilation provides massive speedups for evolving offspring.
# Features
Kozax allows the user to:
- define custom operators
- define custom fitness functions
- use trees flexibly, ranging from symbolic regression to reinforcement learning
- evolve multiple trees simultaneously, even with different inputs
- numerically optimise constants in the computer programs
# How to use
Below is a short demo showing how you can use Kozax. First we generate data:
```python
import jax
import jax.numpy as jnp
import jax.random as jr
key = jr.PRNGKey(0)
key, data_key, init_key = jr.split(key, 3)
x = jr.uniform(data_key, shape=(30,), minval=-5, maxval = 5)
y = -0.1*x**3 + 0.3*x**2 + 1.5*x
```
Now we have to define a fitness function. This allows for much freedom, because you can use the computer program anyway you want to during evaluation. The fitness function should have a `__call__` method that receives a candidate, the data and a function that is necessary to evaluate the tree.
```python
class FitnessFunction:
def __call__(self, candidate, data, tree_evaluator):
_X, _Y = data
pred = jax.vmap(tree_evaluator, in_axes=[None, 0])(candidate, _X)
return jnp.mean(jnp.square(pred-_Y))
fitness_function = FitnessFunction()
```
Now we will use genetic programming to recover the equation from the data. This requires defining the hyperparameters, initializing the population and the general loop of evaluating and evolving the population.
```python
from Kozax.genetic_programming import GeneticProgramming
#Define hyperparameters
population_size = 500
num_generations = 100
strategy = GeneticProgramming(num_generations, population_size, fitness_function)
population = strategy.initialize_population(init_key)
for g in range(num_generations):
key, eval_key, sample_key = jr.split(key, 3)
fitness, population = strategy.evaluate_population(population, (x[:,None], y[:,None]), eval_key)
if g < (num_generations-1):
population = strategy.evolve(population, fitness, sample_key)
best_fitnesses, best_solutions = strategy.get_statistics()
print(f"The best solution is {strategy.to_string(best_solutions[-1])} with a fitness of {best_fitnesses[-1]}")
```
There are additional [examples](https://github.com/sdevries0/Kozax/tree/main/examples) on how to use Kozax on more complex problems.
# Citation
If you make use of this code in your research paper, please cite:
```
@article{de2024discovering,
title={Discovering Dynamic Symbolic Policies with Genetic Programming},
author={de Vries, Sigur and Keemink, Sander and van Gerven, Marcel},
journal={arXiv preprint arXiv:2406.02765},
year={2024}
}
```
Raw data
{
"_id": null,
"home_page": null,
"name": "Kozax",
"maintainer": null,
"docs_url": null,
"requires_python": "~=3.10",
"maintainer_email": null,
"keywords": "Evolutionary algorithms, Genetic programming, JAX, Symbolic regression",
"author": null,
"author_email": "Sigur de Vries <sigur@hotmail.nl>",
"download_url": "https://files.pythonhosted.org/packages/58/7c/a1f3fdf9203e01fe0d3a275777d344e667403e0b268e54cefd31ca55663b/kozax-0.0.1.tar.gz",
"platform": null,
"description": "# Kozax: Genetic Programming in JAX\nKozax introduces a general framework for evolving computer programs with genetic programming in JAX. With JAX, the computer programs can be vectorized and evaluated on parallel on CPU and GPU. Furthermore, Just-in-time compilation provides massive speedups for evolving offspring.\n\n# Features\nKozax allows the user to:\n- define custom operators\n- define custom fitness functions\n- use trees flexibly, ranging from symbolic regression to reinforcement learning\n- evolve multiple trees simultaneously, even with different inputs\n- numerically optimise constants in the computer programs\n\n# How to use\nBelow is a short demo showing how you can use Kozax. First we generate data:\n```python\nimport jax\nimport jax.numpy as jnp\nimport jax.random as jr\n\nkey = jr.PRNGKey(0)\nkey, data_key, init_key = jr.split(key, 3)\nx = jr.uniform(data_key, shape=(30,), minval=-5, maxval = 5)\ny = -0.1*x**3 + 0.3*x**2 + 1.5*x\n```\n\nNow we have to define a fitness function. This allows for much freedom, because you can use the computer program anyway you want to during evaluation. The fitness function should have a `__call__` method that receives a candidate, the data and a function that is necessary to evaluate the tree.\n```python\nclass FitnessFunction:\n def __call__(self, candidate, data, tree_evaluator):\n _X, _Y = data\n pred = jax.vmap(tree_evaluator, in_axes=[None, 0])(candidate, _X)\n return jnp.mean(jnp.square(pred-_Y))\n\nfitness_function = FitnessFunction()\n```\n\nNow we will use genetic programming to recover the equation from the data. This requires defining the hyperparameters, initializing the population and the general loop of evaluating and evolving the population.\n```python\nfrom Kozax.genetic_programming import GeneticProgramming\n\n#Define hyperparameters\npopulation_size = 500\nnum_generations = 100\n\nstrategy = GeneticProgramming(num_generations, population_size, fitness_function)\n\npopulation = strategy.initialize_population(init_key)\n\nfor g in range(num_generations):\n key, eval_key, sample_key = jr.split(key, 3)\n fitness, population = strategy.evaluate_population(population, (x[:,None], y[:,None]), eval_key)\n\n if g < (num_generations-1):\n population = strategy.evolve(population, fitness, sample_key)\n\nbest_fitnesses, best_solutions = strategy.get_statistics()\nprint(f\"The best solution is {strategy.to_string(best_solutions[-1])} with a fitness of {best_fitnesses[-1]}\")\n```\n\nThere are additional [examples](https://github.com/sdevries0/Kozax/tree/main/examples) on how to use Kozax on more complex problems.\n\n\n# Citation\nIf you make use of this code in your research paper, please cite:\n```\n@article{de2024discovering,\n title={Discovering Dynamic Symbolic Policies with Genetic Programming},\n author={de Vries, Sigur and Keemink, Sander and van Gerven, Marcel},\n journal={arXiv preprint arXiv:2406.02765},\n year={2024}\n}\n```\n",
"bugtrack_url": null,
"license": "Kozax: Genetic programming framework in JAX Copyright (c) 2024 sdevries0 This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.",
"summary": "Genetic programming library in JAX.",
"version": "0.0.1",
"project_urls": null,
"split_keywords": [
"evolutionary algorithms",
" genetic programming",
" jax",
" symbolic regression"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "18e013d4f55b5ffe5ae5e5077194de206aaa5784ebf3d69daa0933c6c57bdbac",
"md5": "c8ea1d456a2ddd69b34658c02ebce0fd",
"sha256": "88c8550e3393fbad2b6fefb98148ab6a31e2a4acdad3bea6947f074f9eedfeb8"
},
"downloads": -1,
"filename": "kozax-0.0.1-py3-none-any.whl",
"has_sig": false,
"md5_digest": "c8ea1d456a2ddd69b34658c02ebce0fd",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": "~=3.10",
"size": 48444,
"upload_time": "2025-01-10T11:11:20",
"upload_time_iso_8601": "2025-01-10T11:11:20.558611Z",
"url": "https://files.pythonhosted.org/packages/18/e0/13d4f55b5ffe5ae5e5077194de206aaa5784ebf3d69daa0933c6c57bdbac/kozax-0.0.1-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "587ca1f3fdf9203e01fe0d3a275777d344e667403e0b268e54cefd31ca55663b",
"md5": "5797e8303e121fd579f79764c30c75d6",
"sha256": "e8a4dd2b694b51906166309ae2b756ebb9cc96de968495e9417324612a61ca6f"
},
"downloads": -1,
"filename": "kozax-0.0.1.tar.gz",
"has_sig": false,
"md5_digest": "5797e8303e121fd579f79764c30c75d6",
"packagetype": "sdist",
"python_version": "source",
"requires_python": "~=3.10",
"size": 28988,
"upload_time": "2025-01-10T11:11:25",
"upload_time_iso_8601": "2025-01-10T11:11:25.860529Z",
"url": "https://files.pythonhosted.org/packages/58/7c/a1f3fdf9203e01fe0d3a275777d344e667403e0b268e54cefd31ca55663b/kozax-0.0.1.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-01-10 11:11:25",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "kozax"
}
