EvoXTorch project structure changes

.github/pull_request_template.md

@@ -8,7 +8,7 @@ Please describe your pull request here
 A quick checklist, please check what applies to your pull request (put "x" in the brackets)
 -->
 
-- [ ] I have formatted my Python code with `black`.
+- [ ] I have formatted my Python code with `ruff`.
 - [ ] I have good commit messages.
 - If adding new algorithms, problems, operators:
   - [ ] Added related test cases.

.github/workflows/python-package.yml

@@ -29,21 +29,11 @@ jobs:
     - name: Install build dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install build wheel setuptools pytest
-    - name: Build and install package Python 3.10 and above
+        python -m pip install build wheel setuptools
+    - name: Build and install package
       run: |
         output=$(python -m build --wheel)
-        if [[ "${{ matrix.python-version }}" == "3.12" ]]; then
-          # Exclude envpool when testing with Python 3.12
-          pip install dist/${output##* }[test]
-          pip install gymnasium ray gpjax plotly pandas tensorflow-datasets grain brax plotly pandas
-        else
-          pip install dist/${output##* }[full,test]
-        fi
-    - name: Test with pytest Python 3.10 and above
+        pip install dist/${output##* }[test]
+    - name: Test with unittest
       run: |
-        if [[ "${{ matrix.python-version }}" == "3.12" ]]; then
-          pytest -k 'not test_envpool_cartpole'
-        else
-          pytest
-        fi
+        python -m unittest

.github/workflows/ruff-check.yml

@@ -0,0 +1,8 @@
+name: Ruff
+on: [ push, pull_request ]
+jobs:
+  ruff:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/ruff-action@v3

.gitignore

@@ -166,6 +166,5 @@ cython_debug/
 
 # Test
 tests
-evox
 **/_future/*
-data
+data/MNIST

.pre-commit-config.yaml

@@ -0,0 +1,19 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v3.2.0
+  hooks:
+  - id: trailing-whitespace
+  - id: end-of-file-fixer
+  - id: check-yaml
+  - id: check-added-large-files
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.8.5
+  hooks:
+    # Run the linter.
+    - id: ruff
+      args: [ --fix ]
+    # Run the formatter.
+    - id: ruff-format

README.md

@@ -24,7 +24,7 @@
 
 ---
 
-Building upon [JAX](https://github.com/google/jax) and [Ray](https://github.com/ray-project/ray), EvoX offers a comprehensive suite of **50+ Evolutionary Algorithms (EAs)** and a wide range of **100+ Benchmark Problems/Environments**, all benefiting from distributed GPU-acceleration. It facilitates efficient exploration of complex optimization landscapes, effective tackling of black-box optimization challenges, and deep dives into neuroevolution with [Brax](https://github.com/google/brax). With a foundation in functional programming and hierarchical state management, EvoX offers a user-friendly and modular experience. For more details, please refer to our [Paper](https://arxiv.org/abs/2301.12457) and [Documentation](https://evox.readthedocs.io/en/latest/) / [文档](https://evox.readthedocs.io/zh/latest/).
+Building upon PyTorch, EvoX offers a comprehensive suite of **50+ Evolutionary Algorithms (EAs)** and a wide range of **100+ Benchmark Problems/Environments**, all benefiting from distributed GPU-acceleration. It facilitates efficient exploration of complex optimization landscapes, effective tackling of black-box optimization challenges, and deep dives into neuroevolution with [Brax](https://github.com/google/brax). With a foundation in functional programming and hierarchical state management, EvoX offers a user-friendly and modular experience. For more details, please refer to our [Paper](https://arxiv.org/abs/2301.12457) and [Documentation](https://evox.readthedocs.io/en/latest/) / [文档](https://evox.readthedocs.io/zh/latest/).
 
 ## Key Features
 
@@ -63,107 +63,33 @@ For a comprehensive list and further details of all algorithms, please check the
 
 ### Benchmark Problems/Environments
 
-| Category       | Problems/Environments               |
-| -------------- | ----------------------------------- |
-| Numerical      | DTLZ, LSMOP, MaF, ZDT, CEC'22,  ... |
+| Category          | Problems/Environments               |
+| ----------------- | ----------------------------------- |
+| Numerical         | DTLZ, LSMOP, MaF, ZDT, CEC'22,  ... |
 | Neuroevolution/RL | Brax, Gym, TorchVision Dataset, ... |
 
 For a comprehensive list and further details of all benchmark problems/environments, please check the [API Documentation](https://evox.readthedocs.io/en/latest/api/problems/index.html).
 
 
 ## Setting Up EvoX
 
-
-## Prerequisites
-
-- **Python**: Version 3.12 (or higher)
-- **CUDA**: Version 12.1 (or higher)
-- **PyTorch**: Version 2.5.0 (or higher recommended)
-
 Install `evox` effortlessly via `pip`:
 ```bash
 pip install evox
 ```
 
-**Note**: To setup EvoX with **GPU acceleration** capabilities, you will need to setup **JAX** first. For detials, please refer to our comprehensive [Installation Guide](https://evox.readthedocs.io/en/latest/guide/install/index.html). Additionally, you can watch our **instructional videos**:
-
-🎥 [EvoX Installation Guide (Linux)](https://youtu.be/fa2s1Jl-Fy0)
-
-🎥 [EvoX Installation Guide (Windows)](https://youtu.be/7f8Uz1rqvn8)
-
-🎥 [EvoX 安装指南 (Linux)](https://www.bilibili.com/video/BV1Zt421c7GN)
-
-🎥 [EvoX 安装指南 (Windows)](https://www.bilibili.com/video/BV1Bb421h7bG)
-
-
-
-## Quick Start
-
-Kickstart your journey with EvoX in just a few simple steps:
-1. **Import necessary modules**:
-```python
-import evox
-from evox import algorithms, problems, workflows
-```
-2. **Configure an algorithm and define a problem**:
-```python
-pso = algorithms.PSO(
-    lb=jnp.full(shape=(2,), fill_value=-32),
-    ub=jnp.full(shape=(2,), fill_value=32),
-    pop_size=100,
-)
-ackley = problems.numerical.Ackley()
-```
-3. **Compose and initialize the workflow**:
-```python
-workflow = workflows.StdWorkflow(pso, ackley)
-key = jax.random.PRNGKey(42)
-state = workflow.init(key)
-```
-4. **Run the workflow**:
-```python
-# Execute the workflow for 100 iterations
-for i in range(100):
-    state = workflow.step(state)
-```
-
-## Use-cases and Applications
-
-Try out ready-to-play examples in your browser with Colab:
-
-| Example                  | Link                                                                                                                                                                                                   |
-| ------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| Basic Usage              | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EMI-Group/evox/blob/main/docs/source/guide/basics/1-start.ipynb)                 |
-| Numerical Optimization   | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EMI-Group/evox/blob/main/docs/source/example/pso_ackley.ipynb)                   |
-| Neuroevolution with Gym  | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EMI-Group/evox/blob/main/docs/source/example/gym_classic_control.ipynb)          |
-| Neuroevolution with Brax | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EMI-Group/evox/blob/main/docs/source/guide/basics/2-problems.ipynb)              |
-| Custom Algorithm/Problem | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EMI-Group/evox/blob/main/docs/source/example/custom_algorithm_and_problem.ipynb) |
-
-For more use-cases and applications, pleae check out [Example Directory](https://evox.readthedocs.io/en/latest/example/index.html).
-
-
-## Unit Test Commands
+## Run Unit Test
 
+1. prepare the test environment by installing the required packages (e.g., `torch`) in your Python environment
+2. run the following command in the **root directory of the project**: `pip install -e .`, which will install the `evox` package in editable mode
+3. run unittest:
 ```shell
-python ./unit_test/algorithms/pso_variants/test_clpso.py
-python ./unit_test/algorithms/pso_variants/test_cso.py
-python ./unit_test/algorithms/pso_variants/test_dms_pso_el.py
-python ./unit_test/algorithms/pso_variants/test_fs_pso.py
-python ./unit_test/algorithms/pso_variants/test_pso.py
-python ./unit_test/algorithms/pso_variants/test_sl_pso_gs.py
-python ./unit_test/algorithms/pso_variants/test_sl_pso_us.py
-
-python ./unit_test/core/test_jit_util.py 
-python ./unit_test/core/test_module.py 
-
-python ./unit_test/problems/test_hpo_wrapper.py 
-python ./unit_test/problems/test_supervised_learning.py 
-
-python ./unit_test/utils/test_jit_fix.py 
-python ./unit_test/utils/test_parameters_and_vector.py
-python ./unit_test/utils/test_while.py 
-
-python ./unit_test/workflows/test_std_workflow.py
+# run all tests
+python -m unittest
+# run tests in [path], e.g. python -m unittest unit_test/core/test_jit_util.py
+python -m unittest [path-to-test-file]
+# run a specific test method or module, e.g. python -m unittest unit_test.core.test_jit_util.TestJitUtil.test_single_eval
+python -m unittest [path-to-method-or-module]
 ```
 
 ## Community & Support
@@ -195,4 +121,4 @@ If you use EvoX in your research and want to cite it in your work, please use:
 
 ## Star History
 
-[![Star History Chart](https://api.star-history.com/svg?repos=EMI-Group/evox&type=Date)](https://star-history.com/#EMI-Group/evox&Date)
+[![Star History Chart](https://api.star-history.com/svg?repos=EMI-Group/evox&type=Date)](https://star-history.com/#EMI-Group/evox&Date)

benchmarks/pso.py

@@ -0,0 +1,60 @@
+"""Benchmark the performance of PSO algorithm in EvoX."""
+
+import time
+
+import torch
+from torch.profiler import ProfilerActivity, profile
+
+from evox.algorithms import PSO
+from evox.core import Problem, jit, use_state, vmap
+from evox.workflows import StdWorkflow
+
+
+def run_pso():
+    class Sphere(Problem):
+        def __init__(self):
+            super().__init__()
+
+        def evaluate(self, pop: torch.Tensor):
+            return (pop**2).sum(-1)
+
+    algo = PSO(pop_size=10, lb=-10 * torch.ones(3), ub=10 * torch.ones(3))
+    prob = Sphere()
+
+    torch.set_default_device("cuda" if torch.cuda.is_available() else "cpu")
+    print(torch.get_default_device())
+    workflow = StdWorkflow()
+    workflow.setup(algo, prob)
+    workflow.init_step()
+    workflow.step()
+    state_step = use_state(lambda: workflow.step)
+    vmap_state_step = vmap(state_step)
+    print(vmap_state_step.init_state(2))
+    state = state_step.init_state()
+    jit_state_step = jit(state_step, trace=True, example_inputs=(state,))
+    state = state_step.init_state()
+    t = time.time()
+    with profile(
+        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+        record_shapes=True,
+        profile_memory=True,
+    ) as prof:
+        for _ in range(1000):
+            workflow.step()
+    print(prof.key_averages().table())
+    torch.cuda.synchronize()
+    t = time.time()
+    with profile(
+        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+        record_shapes=True,
+        profile_memory=True,
+    ) as prof:
+        for _ in range(1000):
+            state = jit_state_step(state)
+    print(prof.key_averages().table())
+    torch.cuda.synchronize()
+    print(time.time() - t)
+
+
+if __name__ == "__main__":
+    run_pso()

benchmarks/test_base.py

@@ -0,0 +1,92 @@
+import os
+
+import torch
+from torch.profiler import ProfilerActivity, profile
+
+from src.core import Algorithm, Problem, jit, use_state, vmap
+from src.workflows import EvalMonitor, StdWorkflow
+
+
+class Sphere(Problem):
+    def __init__(self):
+        super().__init__()
+
+    def evaluate(self, pop: torch.Tensor):
+        return (pop**2).sum(-1)
+
+
+def test(
+    algo: Algorithm,
+    print_path: str | None = None,
+    profiling: bool = True,
+    test_trace: bool = True,
+):
+    torch.set_default_device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Current device: ", torch.get_default_device())
+
+    monitor = EvalMonitor(full_fit_history=False, full_sol_history=False)
+    prob = Sphere()
+    workflow = StdWorkflow()
+    workflow.setup(algo, prob)
+    # test trace step
+    if test_trace:
+        state_step = use_state(lambda: workflow.step)
+        state = state_step.init_state()
+        jit_state_step = jit(state_step, trace=True, example_inputs=(state,))
+        vmap_state_step = vmap(state_step)
+        vmap_state_step = jit(
+            vmap_state_step,
+            trace=True,
+            lazy=False,
+            example_inputs=(vmap_state_step.init_state(3),),
+        )
+    # print
+    if print_path is not None:
+        with open(os.path.join(print_path, "script.md"), "w") as ff:
+            ff.write(workflow.step.inlined_graph.__str__())
+        if test_trace:
+            with open(os.path.join(print_path, "trace.md"), "w") as ff:
+                ff.write(jit_state_step.inlined_graph.__str__())
+            with open(os.path.join(print_path, "vmap.md"), "w") as ff:
+                ff.write(vmap_state_step.inlined_graph.__str__())
+    # profile
+    workflow = StdWorkflow()
+    workflow.setup(algo, prob, monitor)
+    workflow.init_step()
+    print("Initial best fitness:", workflow.get_submodule("monitor").topk_fitness)
+    if profiling:
+        with profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+            record_shapes=True,
+            profile_memory=True,
+        ) as prof:
+            for _ in range(1000):
+                workflow.step()
+        print(prof.key_averages().table())
+    else:
+        for _ in range(1000):
+            workflow.step()
+    print("Final best fitness:", workflow.get_submodule("monitor").topk_fitness)
+    torch.cuda.synchronize()
+    if test_trace:
+        workflow = StdWorkflow()
+        workflow.setup(algo, prob, monitor)
+        workflow.init_step()
+        state_step = use_state(lambda: workflow.step)
+        state = state_step.init_state()
+        print("Initial best fitness:", state["self.algorithm._monitor_.topk_fitness"])
+        jit_state_step = jit(state_step, trace=True, example_inputs=(state,))
+        if profiling:
+            with profile(
+                activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+                record_shapes=True,
+                profile_memory=True,
+            ) as prof:
+                for _ in range(1000):
+                    state = jit_state_step(state)
+            print(prof.key_averages().table())
+        else:
+            for _ in range(1000):
+                state = jit_state_step(state)
+        print("Final best fitness:", state["self.algorithm._monitor_.topk_fitness"])
+        torch.cuda.synchronize()

benchmarks/utils.py

@@ -0,0 +1,32 @@
+"""Benchmark the performance of utils functions in EvoX"""
+
+import torch
+from torch.profiler import ProfilerActivity, profile
+
+from evox.core import jit
+from evox.utils import switch
+
+
+def run_switch():
+    x = torch.tensor([1, 0, 1], dtype=torch.int)
+    y = torch.tensor([[2.0, 2.5], [3.0, 3.5], [4.0, 4.5]]).T.split(1, dim=0)
+    y = [a.squeeze(0) for a in y]
+    basic_switch = jit(switch, trace=False)
+    z = basic_switch(x, y)
+    print(z)
+
+    x = torch.randint(low=0, high=10, size=(1000, 10000), dtype=torch.int, device="cuda")
+    y = [torch.rand(1000, 10000, device="cuda") for _ in range(10)]
+    vmap_switch = jit(switch, trace=False)
+    with profile(
+        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+        record_shapes=True,
+        profile_memory=True,
+    ) as prof:
+        for _ in range(1000):
+            z = vmap_switch(x, y)
+    print(prof.key_averages().table())
+
+
+if __name__ == "__main__":
+    run_switch()