├── .gitignore ├── LICENSE ├── README.md ├── firefly.py ├── firefly_with_colony_resets.py └── requirements.txt /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control 110 | .pdm.toml 111 | .pdm-python 112 | .pdm-build/ 113 | 114 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 115 | __pypackages__/ 116 | 117 | # Celery stuff 118 | celerybeat-schedule 119 | celerybeat.pid 120 | 121 | # SageMath parsed files 122 | *.sage.py 123 | 124 | # Environments 125 | .env 126 | .venv 127 | env/ 128 | venv/ 129 | ENV/ 130 | env.bak/ 131 | venv.bak/ 132 | 133 | # Spyder project settings 134 | .spyderproject 135 | .spyproject 136 | 137 | # Rope project settings 138 | .ropeproject 139 | 140 | # mkdocs documentation 141 | /site 142 | 143 | # mypy 144 | .mypy_cache/ 145 | .dmypy.json 146 | dmypy.json 147 | 148 | # Pyre type checker 149 | .pyre/ 150 | 151 | # pytype static type analyzer 152 | .pytype/ 153 | 154 | # Cython debug symbols 155 | cython_debug/ 156 | 157 | # PyCharm 158 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 159 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 160 | # and can be added to the global gitignore or merged into this file. For a more nuclear 161 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 162 | #.idea/ 163 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Phil Wang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Firefly Algorithm - Pytorch 2 | 3 | Exploration into the Firefly algorithm (a generalized version of particle swarm optimization) in Pytorch. In particular interested in hybrid firefly + genetic algorithms, or ones that are gender-based. 4 | 5 | ## Install 6 | 7 | ```bash 8 | $ pip install -r requirements.txt 9 | ``` 10 | 11 | ## Usage 12 | 13 | Test on rosenbrock minimization 14 | 15 | ```bash 16 | $ python firefly.py --use-genetic-algorithm 17 | ``` 18 | 19 | ## Citations 20 | 21 | ```bibtex 22 | @article{Yang2018WhyTF, 23 | title = {Why the Firefly Algorithm Works?}, 24 | author = {Xin-She Yang and Xingshi He}, 25 | journal = {ArXiv}, 26 | year = {2018}, 27 | volume = {abs/1806.01632}, 28 | url = {https://api.semanticscholar.org/CorpusID:46940737} 29 | } 30 | ``` 31 | 32 | ```bibtex 33 | @article{article, 34 | author = {El-Shorbagy, M. and Elrefaey, Adel}, 35 | year = {2022}, 36 | month = {04}, 37 | pages = {706-730}, 38 | title = {A hybrid genetic-firefly algorithm for engineering design problems}, 39 | volume = {Journal of Computational Design and Engineering, Volume 9}, 40 | journal = {Journal of Computational Design and Engineering}, 41 | doi = {10.1093/jcde/qwac013} 42 | } 43 | ``` 44 | -------------------------------------------------------------------------------- /firefly.py: -------------------------------------------------------------------------------- 1 | import fire 2 | import torch 3 | import einx # s - colonies, p - population, i - population source, j - population target, t - tournament participants, d - dimension 4 | 5 | # test objective function - solution is close to all 1.'s 6 | 7 | def rosenbrock(x): 8 | return (100 * (x[..., 1:] - x[..., :-1] ** 2) ** 2 + (1 - x[..., :-1]) ** 2).sum(dim = -1) 9 | 10 | # hyperparameters 11 | 12 | @torch.inference_mode() 13 | def main( 14 | steps = 5000, 15 | colonies = 4, 16 | population_size = 1000, 17 | dimensions = 12, # set this to something lower (2-10) for fireflies without sexual reproduction to solve 18 | lower_bound = -4., 19 | upper_bound = 4., 20 | migrate_every = 50, 21 | 22 | beta0 = 2., # exploitation factor, moving fireflies of low light intensity to high 23 | gamma = 1., # controls light intensity decay over distance - setting this to zero will make firefly equivalent to vanilla PSO 24 | alpha = 0.1, # exploration factor 25 | alpha_decay = 0.995, # exploration decay each step 26 | 27 | # genetic algorithm related 28 | 29 | use_genetic_algorithm = False, # turn on genetic algorithm, for comparing with non-sexual fireflies 30 | breed_every = 10, 31 | tournament_size = 50, 32 | num_children = 250 33 | ): 34 | 35 | assert tournament_size <= population_size 36 | assert num_children <= population_size 37 | 38 | # settings 39 | 40 | use_cuda = True 41 | verbose = True 42 | 43 | cost_function = rosenbrock 44 | 45 | # main algorithm 46 | 47 | fireflies = torch.zeros((colonies, population_size, dimensions)).uniform_(lower_bound, upper_bound) 48 | 49 | # maybe use cuda 50 | 51 | if torch.cuda.is_available() and use_cuda: 52 | fireflies = fireflies.cuda() 53 | 54 | device = fireflies.device 55 | 56 | # iterate 57 | 58 | for step in range(steps): 59 | 60 | # cost, which is inverse of light intensity 61 | 62 | costs = cost_function(fireflies) 63 | 64 | if verbose: 65 | print(f'{step}: {costs.amin():.5f}') 66 | 67 | # fireflies with lower light intensity (high cost) moves towards the higher intensity (lower cost) 68 | 69 | move_mask = einx.greater('s i, s j -> s i j', costs, costs) 70 | 71 | # get vectors of fireflies to one another 72 | # calculate distance and the beta 73 | 74 | delta_positions = einx.subtract('s j d, s i d -> s i j d', fireflies, fireflies) 75 | 76 | distance = delta_positions.norm(dim = -1) 77 | 78 | betas = beta0 * (-gamma * distance ** 2).exp() 79 | 80 | # calculate movements 81 | 82 | attraction = einx.multiply('s i j, s i j d -> s i j d', move_mask * betas, delta_positions) 83 | random_walk = alpha * (torch.rand_like(fireflies) - 0.5) * (upper_bound - lower_bound) 84 | 85 | # move the fireflies 86 | 87 | fireflies += einx.sum('s i j d -> s i d', attraction) + random_walk 88 | 89 | fireflies.clamp_(min = lower_bound, max = upper_bound) 90 | 91 | # decay exploration factor 92 | 93 | alpha *= alpha_decay 94 | 95 | # have colonies migrate every so often 96 | 97 | if colonies > 1 and migrate_every > 0 and (step % migrate_every) == 0: 98 | midpoint = population_size // 2 99 | fireflies, fireflies_rotate = fireflies[:, :midpoint], fireflies[:, midpoint:] 100 | migrate_indices = torch.randperm(colonies, device = device) 101 | fireflies = torch.cat((fireflies, fireflies_rotate[migrate_indices]), dim = 1) 102 | 103 | # maybe genetic algorithm 104 | 105 | if not use_genetic_algorithm or (step % breed_every) != 0: 106 | continue 107 | 108 | # use the most effective genetic algorithm - tournament style 109 | 110 | cost = cost_function(fireflies) 111 | fitness = 1. / cost 112 | 113 | batch_randperm = torch.randn((colonies, num_children, population_size), device = device).argsort(dim = -1) 114 | tournament_indices = batch_randperm[..., :tournament_size] 115 | 116 | participant_fitnesses = einx.get_at('s [p], s c t -> s c t', fitness, tournament_indices) 117 | winner_tournament_ids = participant_fitnesses.topk(2, dim = -1).indices 118 | 119 | winning_firefly_indices = einx.get_at('s c [t], s c parents -> s c parents', tournament_indices, winner_tournament_ids) 120 | 121 | # breed the top two winners of each tournament 122 | 123 | parent1, parent2 = einx.get_at('s [p] d, s c parents -> parents s c d', fireflies, winning_firefly_indices) 124 | 125 | # do a uniform crossover 126 | 127 | crossover_mask = torch.rand_like(parent1) < 0.5 128 | 129 | children = torch.where(crossover_mask, parent1, parent2) 130 | 131 | # sort the fireflies by fitness and replace the worst performing with the new children 132 | 133 | replacement_mask = fitness.argsort(dim = -1).argsort(dim = -1) < num_children 134 | 135 | fireflies[replacement_mask] = einx.rearrange('s p d -> (s p) d', children) 136 | 137 | # print solution 138 | 139 | fireflies = einx.rearrange('s p d -> (s p) d', fireflies) 140 | 141 | costs = cost_function(fireflies) 142 | sorted_costs, sorted_indices = costs.sort(dim = -1) 143 | 144 | fireflies = fireflies[sorted_indices] 145 | 146 | print(f'best performing firefly for rosenbrock with {dimensions} dimensions: {sorted_costs[0]:.3f}: {fireflies[0]}') 147 | 148 | # main 149 | 150 | if __name__ == '__main__': 151 | fire.Fire(main) 152 | -------------------------------------------------------------------------------- /firefly_with_colony_resets.py: -------------------------------------------------------------------------------- 1 | import fire 2 | import torch 3 | import einx # s - colonies, p - population, i - population source, j - population target, t - tournament participants, d - dimension 4 | 5 | # test objective function - solution is close to all 1.'s 6 | 7 | def rosenbrock(x): 8 | return (100 * (x[..., 1:] - x[..., :-1] ** 2) ** 2 + (1 - x[..., :-1]) ** 2).sum(dim = -1) 9 | 10 | # genetic algorithm functions 11 | 12 | def children_from_tournaments( 13 | fireflies, 14 | fitness, 15 | num_children, 16 | tournament_size 17 | ): 18 | device = fireflies.device 19 | 20 | shape = list(fireflies.shape) 21 | shape[-2] = num_children 22 | 23 | batch_randperm = torch.randn(shape, device = device).argsort(dim = -1) 24 | tournament_indices = batch_randperm[..., :tournament_size] 25 | 26 | participant_fitnesses = einx.get_at('... [p], ... c t -> ... c t', fitness, tournament_indices) 27 | winner_tournament_ids = participant_fitnesses.topk(2, dim = -1).indices 28 | 29 | winning_firefly_indices = einx.get_at('... c [t], ... c parents -> ... c parents', tournament_indices, winner_tournament_ids) 30 | 31 | # breed the top two winners of each tournament 32 | 33 | parent1, parent2 = einx.get_at('... [p] d, ... c parents -> parents ... c d', fireflies, winning_firefly_indices) 34 | 35 | # do a uniform crossover 36 | 37 | crossover_mask = torch.rand_like(parent1) < 0.5 38 | 39 | children = torch.where(crossover_mask, parent1, parent2) 40 | return children 41 | 42 | # hyperparameters 43 | 44 | @torch.inference_mode() 45 | def main( 46 | steps = 5000, 47 | colonies = 5, 48 | population_size = 1000, 49 | dimensions = 15, # set this to something lower (2-10) for fireflies without sexual reproduction to solve 50 | lower_bound = -4., 51 | upper_bound = 4., 52 | migrate_every = 200, 53 | frac_migrate = 0.1, 54 | 55 | beta0 = 2., # exploitation factor, moving fireflies of low light intensity to high 56 | gamma = 1., # controls light intensity decay over distance - setting this to zero will make firefly equivalent to vanilla PSO 57 | alpha = 0.1, # exploration factor 58 | alpha_decay = 0.995, # exploration decay each step 59 | 60 | # genetic algorithm related 61 | 62 | use_genetic_algorithm = False, # turn on genetic algorithm, for comparing with non-sexual fireflies 63 | breed_every = 5, 64 | tournament_size = 100, 65 | num_children = 250, 66 | 67 | # colonies / island resets 68 | 69 | num_colonies_reset = 1, 70 | reset_colonies_every = 400, 71 | reset_frac = 0.95, # need to preserve some of the elite performers 72 | reset_tournament_size = 25, # when resetting, go for more diversity 73 | ): 74 | 75 | assert tournament_size <= population_size 76 | assert num_children <= population_size 77 | 78 | # settings 79 | 80 | use_cuda = True 81 | verbose = True 82 | 83 | cost_function = rosenbrock 84 | 85 | # main algorithm 86 | 87 | fireflies = torch.zeros((colonies, population_size, dimensions)).uniform_(lower_bound, upper_bound) 88 | 89 | # maybe use cuda 90 | 91 | if torch.cuda.is_available() and use_cuda: 92 | fireflies = fireflies.cuda() 93 | 94 | device = fireflies.device 95 | 96 | # iterate 97 | 98 | for step in range(steps): 99 | 100 | # cost, which is inverse of light intensity 101 | 102 | costs = cost_function(fireflies) 103 | 104 | if verbose: 105 | print(f'{step}: {costs.amin():.5f}') 106 | 107 | # fireflies with lower light intensity (high cost) moves towards the higher intensity (lower cost) 108 | 109 | move_mask = einx.greater('s i, s j -> s i j', costs, costs) 110 | 111 | # get vectors of fireflies to one another 112 | # calculate distance and the beta 113 | 114 | delta_positions = einx.subtract('s j d, s i d -> s i j d', fireflies, fireflies) 115 | 116 | distance = delta_positions.norm(dim = -1) 117 | 118 | betas = beta0 * (-gamma * distance ** 2).exp() 119 | 120 | # calculate movements 121 | 122 | attraction = einx.multiply('s i j, s i j d -> s i j d', move_mask * betas, delta_positions) 123 | random_walk = alpha * (torch.rand_like(fireflies) - 0.5) * (upper_bound - lower_bound) 124 | 125 | # move the fireflies 126 | 127 | fireflies += einx.sum('s i j d -> s i d', attraction) + random_walk 128 | 129 | fireflies.clamp_(min = lower_bound, max = upper_bound) 130 | 131 | # decay exploration factor 132 | 133 | alpha *= alpha_decay 134 | 135 | # have colonies migrate every so often 136 | 137 | if step > 0 and colonies > 1 and migrate_every > 0 and (step % migrate_every) == 0: 138 | num_migrate = int(population_size * frac_migrate) 139 | fireflies, fireflies_rotate = fireflies[:, :-num_migrate], fireflies[:, -num_migrate:] 140 | fireflies_rotate = torch.roll(fireflies_rotate, 1, dims = (0,)) 141 | fireflies = torch.cat((fireflies, fireflies_rotate), dim = 1) 142 | 143 | # maybe genetic algorithm 144 | 145 | if not use_genetic_algorithm or (step % breed_every) != 0: 146 | continue 147 | 148 | # use the most effective genetic algorithm - tournament style 149 | 150 | cost = cost_function(fireflies) 151 | fitness = 1. / cost 152 | 153 | children = children_from_tournaments(fireflies, fitness, num_children, tournament_size) 154 | 155 | # sort the fireflies by fitness and replace the worst performing with the new children 156 | 157 | replacement_mask = fitness.argsort(dim = -1).argsort(dim = -1) < num_children 158 | 159 | fireflies[replacement_mask] = einx.rearrange('s p d -> (s p) d', children) 160 | 161 | if step > 0 and 0 < num_colonies_reset < colonies and (step % reset_colonies_every) == 0: 162 | cost = cost_function(fireflies) 163 | fitness = 1. / cost 164 | 165 | # determine lowest scoring colonie for reset 166 | 167 | fitness_colonies = fitness.mean(dim = -1) 168 | 169 | sorted_indices = fitness_colonies.sort(dim = -1).indices 170 | 171 | # keep only the best performing colonies 172 | 173 | reset_colonies_indices = sorted_indices[:num_colonies_reset] 174 | top_colonies_indices = sorted_indices[num_colonies_reset:] 175 | 176 | top_colony_fireflies = fireflies[top_colonies_indices] 177 | top_colony_fitness = fitness[top_colonies_indices] 178 | 179 | # repopulate new colonies from the aggregated pool of fireflies 180 | 181 | pooled_fireflies = einx.rearrange('s p d -> (s p) d', top_colony_fireflies) 182 | pooled_fitness = einx.rearrange('s p -> (s p)', top_colony_fitness) 183 | 184 | num_children = int(reset_frac * population_size) 185 | children = children_from_tournaments(pooled_fireflies, pooled_fitness, num_children * num_colonies_reset, reset_tournament_size) 186 | 187 | reset_colony_fireflies = fireflies[reset_colonies_indices] 188 | reset_colony_fitness = fitness[reset_colonies_indices] 189 | 190 | replacement_mask = reset_colony_fitness.argsort(dim = -1).argsort(dim = -1) < num_children 191 | reset_colony_fireflies[replacement_mask] = children 192 | 193 | fireflies = torch.cat((top_colony_fireflies, reset_colony_fireflies), dim = 0) 194 | 195 | # print solution 196 | 197 | fireflies = einx.rearrange('s p d -> (s p) d', fireflies) 198 | 199 | costs = cost_function(fireflies) 200 | sorted_costs, sorted_indices = costs.sort(dim = -1) 201 | 202 | fireflies = fireflies[sorted_indices] 203 | 204 | print(f'best performing firefly for rosenbrock with {dimensions} dimensions: {sorted_costs[0]:.3f}: {fireflies[0]}') 205 | 206 | # main 207 | 208 | if __name__ == '__main__': 209 | fire.Fire(main) 210 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | einx 2 | fire 3 | torch 4 | --------------------------------------------------------------------------------