Metadata-Version: 2.4
Name: kozax
Version: 0.0.8
Summary: Genetic programming library in JAX.
Author-email: Sigur de Vries <sigur@hotmail.nl>
License: kozax: Genetic programming framework in JAX
        
        Copyright (c) 2024 sdevries0
        
        This work is licensed under the Creative Commons Attribution-NonCommercial-NoDerivs 4.0 International License.
        
        You are free to:
        Share — copy and redistribute the material in any medium or format
        The licensor cannot revoke these freedoms as long as you follow the license terms.
        
        Under the following terms:
        Attribution — You must give appropriate credit , provide a link to the license, and indicate if changes were made . You may do so in any reasonable manner, but not in any way that suggests the licensor endorses you or your use.
        NonCommercial — You may not use the material for commercial purposes .
        NoDerivatives — If you remix, transform, or build upon the material, you may not distribute the modified material.
        No additional restrictions — You may not apply legal terms or technological measures that legally restrict others from doing anything the license permits.
        
        Notices:
        You do not have to comply with the license for elements of the material in the public domain or where your use is permitted by an applicable exception or limitation .
        No warranties are given. The license may not give you all of the permissions necessary for your intended use. For example, other rights such as publicity, privacy, or moral rights may limit how you use the material.
        
        License Details:
        This work is licensed under the Creative Commons Attribution-NonCommercial-NoDerivs 4.0 International License. To view a copy of this license, visit https://creativecommons.org/licenses/by-nc-nd/4.0/
License-File: LICENSE.md
Keywords: Evolutionary algorithms,Genetic programming,JAX,Symbolic regression
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Natural Language :: English
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: jax>=0.4.36
Requires-Dist: optax>=0.2.3
Requires-Dist: sympy>=1.13.3
Description-Content-Type: text/markdown

<div align="center">
  <img src="figures/logo.png" width="300">
</div>

# Kozax: Flexible and Scalable 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. Check out the [paper](https://arxiv.org/abs/2502.03047) introducing Kozax.

# 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
You can install Kozax via pip with
```
pip install kozax
```

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) #Initialize key
data_key, gp_key = jr.split(key) #Split key for data and genetic programming
x = jr.uniform(data_key, shape=(30,), minval=-5, maxval = 5) #Inputs
y = -0.1*x**3 + 0.3*x**2 + 1.5*x #Targets
```

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. 
```python
from kozax.fitness_functions.base_fitness_function import BaseFitnessFunction

class FitnessFunction(BaseFitnessFunction):
  """
    The fitness function inherits the class BaseFitnessFunction and should implement the __call__ function, with the candidate, data and tree_evaluator as inputs. The tree_evaluator is used to compute the value of the candidate for each input. jax.vmap is used to vectorize the evaluation of the candidate over the inputs. The candidate's predictions are used to compute the fitness value with the mean squared error.
  """
  def __call__(self, candidate, data, tree_evaluator):
    X, Y = data
    predictions = jax.vmap(tree_evaluator, in_axes=[None, 0])(candidate, X)
    return jnp.mean(jnp.square(predictions-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

#Initialize genetic programming strategy
strategy = GeneticProgramming(num_generations, population_size, fitness_function)

#Fit the strategy on the data. With verbose, we can print the intermediate solutions.
strategy.fit(gp_key, (x, y), verbose = True)
```

There are additional [examples](https://github.com/sdevries0/kozax/tree/main/examples) on how to use kozax on more complex problems.

|Example|Notebook|Script|
|---|---|---|
|Symbolic regression of a dynamical system|[Notebook](examples/example_notebooks/symbolic_regression_dynamical_system.ipynb)|[Script](examples/example_scripts/symbolic_regression_dynamical_system.py)|
|Control policy optimization in Gymnax environment|[Notebook](examples/example_notebooks/control_policy_optimization.ipynb)|[Script](examples/example_scripts/control_policy_optimization.py)|
|Control policy optimization with dynamic memory|[Notebook](examples/example_notebooks/control_policy_optimization_with_memory.ipynb)|[Script](examples/example_scripts/control_policy_optimization_with_memory.py)|
|Optimization of a loss function to train a neural network|[Notebook](examples/example_notebooks/objective_function_optimization.ipynb)|[Script](examples/example_scripts/objective_function_optimization.py)|


# Citation
If you make use of this code in your research paper, please cite:
```
@article{de2025kozax,
  title={Kozax: Flexible and Scalable Genetic Programming in JAX},
  author={de Vries, Sigur and Keemink, Sander W and van Gerven, Marcel AJ},
  journal={arXiv preprint arXiv:2502.03047},
  year={2025}
}
```
