├── src ├── __init__.py ├── swarm │ ├── __init__.py │ ├── optimizers │ │ ├── __init__.py │ │ ├── gd.py │ │ ├── fss.py │ │ ├── pso.py │ │ └── pso_adam.py │ └── utils.py └── setup.py ├── examples ├── images │ ├── ac.gif │ ├── burgers.gif │ ├── burgers_train.png │ └── poisson_train.png ├── PINN │ ├── allen-cahn │ │ ├── AC.pckl │ │ └── ac.py │ ├── poisson │ │ ├── poisson.py │ │ └── poisson_ani.py │ ├── advection │ │ └── advection.py │ ├── diffusion │ │ └── diffusion.py │ └── burgers │ │ └── burgers.py └── regression │ ├── regression.py │ └── regression_ani.py ├── requirements ├── LICENSE ├── .gitignore └── README.md /src/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/swarm/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/images/ac.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/caio-davi/PSO-PINN/HEAD/examples/images/ac.gif -------------------------------------------------------------------------------- /examples/images/burgers.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/caio-davi/PSO-PINN/HEAD/examples/images/burgers.gif -------------------------------------------------------------------------------- /examples/PINN/allen-cahn/AC.pckl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/caio-davi/PSO-PINN/HEAD/examples/PINN/allen-cahn/AC.pckl -------------------------------------------------------------------------------- /examples/images/burgers_train.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/caio-davi/PSO-PINN/HEAD/examples/images/burgers_train.png -------------------------------------------------------------------------------- /examples/images/poisson_train.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/caio-davi/PSO-PINN/HEAD/examples/images/poisson_train.png -------------------------------------------------------------------------------- /requirements: -------------------------------------------------------------------------------- 1 | tensorflow==2.6 2 | numpy==1.19.2 3 | matplotlib==3.3.4 4 | pydoe==0.3.8 5 | 6 | 7 | tensorflow_probability -------------------------------------------------------------------------------- /src/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup(name="swarm", version="0.1", packages=find_packages()) 4 | -------------------------------------------------------------------------------- /src/swarm/optimizers/__init__.py: -------------------------------------------------------------------------------- 1 | import swarm.optimizers.gd as gd 2 | import swarm.optimizers.pso_adam as pso_adam 3 | import swarm.optimizers.pso as pso 4 | 5 | get = { 6 | "gd": gd.gd, 7 | "pso": pso.pso, 8 | "pso_adam": pso_adam.pso, 9 | } 10 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 caio-davi 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. -------------------------------------------------------------------------------- /examples/regression/regression.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import matplotlib.pyplot as plt 3 | import time 4 | import numpy as np 5 | import math 6 | 7 | from swarm.optimizers.pso import pso 8 | from swarm import utils 9 | 10 | np.random.seed(123456) 11 | tf.random.set_seed(123456) 12 | 13 | 14 | # Parameters 15 | layers = [1] + 3 * [5] + [1] 16 | pop_size = 100 17 | n_iter = 2000 18 | x_min = -1 19 | x_max = 1 20 | sample_size = 512 21 | noise = 0.0 22 | 23 | 24 | def objective(x, noise=0): 25 | return tf.cos(math.pi * x) - x 26 | 27 | def get_loss(X, y): 28 | def _loss(w, b): 29 | with tf.GradientTape() as tape: 30 | tape.watch(w) 31 | tape.watch(b) 32 | pred = utils.multilayer_perceptron(w, b, X) 33 | loss = tf.reduce_mean((y - pred) ** 2) 34 | trainable_variables = w + b 35 | grads = tape.gradient(loss, trainable_variables) 36 | return loss, grads 37 | 38 | return _loss 39 | 40 | 41 | X = tf.reshape( 42 | tf.Variable(np.linspace(x_min, x_max, sample_size), dtype="float32"), 43 | [sample_size, 1], 44 | ) 45 | y = objective(X, noise) 46 | 47 | y_min, y_max = tf.math.reduce_min(y), tf.math.reduce_max(y) 48 | 49 | X, y = utils.normalize(X, [x_min, x_max]), utils.normalize(y, [y_min, y_max]) 50 | 51 | opt = pso( 52 | get_loss(X, y), 53 | layers, 54 | n_iter, 55 | pop_size, 56 | 0.9, 57 | 0.8, 58 | 0.5, 59 | gd_alpha=1e-4, 60 | x_min=x_min, 61 | x_max=x_max, 62 | verbose=True, 63 | ) 64 | 65 | start = time.time() 66 | opt.train() 67 | end = time.time() 68 | print("\nTime elapsed: ", end - start) 69 | 70 | nn_w, nn_b = opt.get_best() 71 | 72 | pred = utils.multilayer_perceptron(nn_w, nn_b, X) 73 | 74 | print("L2 error: ", tf.reduce_mean(tf.pow(y - pred, 2)).numpy()) 75 | 76 | plt.plot(tf.squeeze(X), tf.squeeze(y), label="Original Function") 77 | plt.plot(tf.squeeze(X), tf.squeeze(pred), "--", label="Swarm") 78 | plt.legend() 79 | plt.show() 80 | -------------------------------------------------------------------------------- /.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 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 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 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # MacOS annoying folder configs 132 | *.DS_Store 133 | 134 | # vscode configs: 135 | *.vscode 136 | 137 | # Images and videos 138 | *.gif 139 | *.png 140 | *.mp4 141 | !examples/images/* 142 | 143 | # Logs 144 | *.csv -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Physics-Informed Neural Networks Trained with Particle Swarm Optimization 2 | 3 | _Physics-informed neural networks (PINNs) have recently emerged as a promising application of deep learning in a wide range of engineering and scientific problems based on partial differential equation models. However, evidence shows that PINN training by gradient descent displays pathologies and stiffness in gradient flow dynamics. In this paper, we propose the use of a hybrid particle swarm optimization and gradient descent approach to train PINNs. The resulting PSO-PINN algorithm not only mitigates the undesired behaviors of PINNs trained with standard gradient descent, but also presents an ensemble approach to PINN that affords the possibility of robust predictions with quantified uncertainty. Experimental results using the Poisson, advection, and Burgers equations show that PSO-PINN consistently outperforms a baseline PINN trained with Adam gradient descent._ 4 | 5 | The full paper is available [here](https://arxiv.org/pdf/2202.01943.pdf). 6 | 7 | > #### All contents in this repo is _WORK IN PROGRESS_. We may have major updates on the code. 8 | 9 | ## [Getting Start](#setup-virtual-environment) 10 | 11 | ## Some preliminary results: 12 | 13 | ### Allen-Cahn Equation: 14 | 15 | ![Allen-Cahn](examples/images/ac.gif) 16 | 17 | The Allen-Cahn equation is a notary example where gradient descent-based optimizers fail to find a good solution. In this example, even using 10 data points (10 space points over time) it has a hard time finding the optimal solution. The PSO-PINN can alleviate this, smoothing the negative influence of the gradient-based model. Also, notice the variance increasing where the solution is not quite correct, representing the algorithm uncertainty on these points. 18 | 19 | ### Burgers Equation: 20 | 21 | 22 | 23 | The final result is on the left image and the training on the right. Each column on the right image represents a snapshot of the training. The first column has 2k iterations, the second column 4k iterations, and the last one 6k iterations. Note how the variance (uncertainty) decreases over the iterations. 24 | 25 | ### Poisson Equation: 26 | 27 | ![Poisson Equation](examples/images/poisson_train.png) 28 | 29 | Notice how the ensemble converges over the iterations. After a few iterations, we have a "collective agreement", which means that most of the swarm agrees with the solution. The variance can be interpreted as a measure of the ensemble's uncertainty in relation to the given solution. 30 | 31 | ## Setup virtual environment 32 | 33 | Create environment: 34 | 35 | ```bash 36 | python3 -m venv --system-site-packages ./.venv 37 | source ./.venv/bin/activate 38 | ``` 39 | 40 | Install requirements: 41 | 42 | ``` 43 | pip install -r requirements 44 | ``` 45 | 46 | Install the project module (on the root folder): 47 | 48 | ```bash 49 | pip install -e ./src 50 | ``` 51 | 52 | ## Citation 53 | 54 | ``` 55 | @article{davi2022pso, 56 | title={PSO-PINN: Physics-Informed Neural Networks Trained with Particle Swarm Optimization}, 57 | author={Davi, Caio and Braga-Neto, Ulisses}, 58 | journal={arXiv preprint arXiv:2202.01943}, 59 | year={2022} 60 | } 61 | ``` 62 | -------------------------------------------------------------------------------- /examples/regression/regression_ani.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import matplotlib.pyplot as plt 3 | import matplotlib.animation as animation 4 | import numpy as np 5 | import math 6 | 7 | from swarm.optimizers.pso import pso 8 | 9 | from swarm import utils 10 | 11 | 12 | np.random.seed(1234) 13 | tf.random.set_seed(1234) 14 | 15 | # Parameters 16 | layer_sizes = [1] + 3 * [5] + [1] 17 | pop_size = 100 18 | n_iter = 20 19 | stepInd = 0.01 20 | stepVol = 0.01 21 | w_scale = 100 22 | x_min = -1 23 | x_max = 1 24 | sample_size = 265 25 | noise = 0.0 26 | 27 | 28 | def objective(x, noise=0): 29 | return (x ** 3 - tf.sin(2 * math.pi * x)) + tf.random.normal( 30 | x.shape, 0, noise, tf.float32 31 | ) 32 | 33 | 34 | def get_loss(X, y): 35 | def _loss(w, b): 36 | with tf.GradientTape() as tape: 37 | tape.watch(w) 38 | tape.watch(b) 39 | pred = utils.multilayer_perceptron(w, b, X) 40 | loss = tf.reduce_mean((y - pred) ** 2) 41 | trainable_variables = w + b 42 | grads = tape.gradient(loss, trainable_variables) 43 | return loss, grads 44 | 45 | return _loss 46 | 47 | 48 | X = tf.reshape( 49 | tf.Variable(np.linspace(x_min, x_max, sample_size), dtype="float32"), 50 | [sample_size, 1], 51 | ) 52 | y = objective(X, noise) 53 | 54 | 55 | opt_pso = pso( 56 | get_loss(X, y), 57 | layer_sizes, 58 | n_iter, 59 | pop_size, 60 | 0.9, 61 | 0.8, 62 | 0.5, 63 | x_min=x_min, 64 | x_max=x_max, 65 | gd_alpha=0, 66 | ) 67 | 68 | opt_pso_gd = pso( 69 | get_loss(X, y), 70 | layer_sizes, 71 | n_iter, 72 | pop_size, 73 | 0.9, 74 | 0.8, 75 | 0.5, 76 | x_min=x_min, 77 | x_max=x_max, 78 | gd_alpha=1e-2, 79 | ) 80 | 81 | opt_pso_gd2 = pso( 82 | get_loss(X, y), 83 | layer_sizes, 84 | n_iter, 85 | pop_size, 86 | 0.9, 87 | 0.8, 88 | 0.5, 89 | x_min=x_min, 90 | x_max=x_max, 91 | gd_alpha=1e-4, 92 | ) 93 | 94 | 95 | s = np.linspace(x_min, x_max, 100) 96 | tests = tf.reshape(tf.Variable(s, dtype="float32"), [100, 1]) 97 | results = objective(tests) 98 | 99 | 100 | def snapshot(i): 101 | 102 | opt_pso.train() 103 | nn_w_pso, nn_b_pso = opt_pso.get_best() 104 | pred_pso = utils.multilayer_perceptron(nn_w_pso, nn_b_pso, tests) 105 | 106 | opt_pso_gd.train() 107 | nn_w_pso_gd, nn_b_pso_gd = opt_pso_gd.get_best() 108 | pred_pso_gd = utils.multilayer_perceptron(nn_w_pso_gd, nn_b_pso_gd, tests) 109 | 110 | opt_pso_gd2.train() 111 | nn_w_pso_gd2, nn_b_pso_gd2 = opt_pso_gd2.get_best() 112 | pred_pso_gd2 = utils.multilayer_perceptron(nn_w_pso_gd2, nn_b_pso_gd2, tests) 113 | 114 | plt.clf() 115 | plt.xlim([-1, 1]) 116 | plt.ylim([-1.5, 1.5]) 117 | 118 | plt.plot( 119 | tf.squeeze(tests), tf.squeeze(results), label="$f(x) = x^3 - \sin (2 \pi x)$" 120 | ) 121 | plt.plot(tf.squeeze(tests), tf.squeeze(pred_pso), "--", label="PSO") 122 | plt.plot( 123 | tf.squeeze(tests), 124 | tf.squeeze(pred_pso_gd), 125 | "--", 126 | label="PSO - GD (lr = 1e-2)", 127 | ) 128 | plt.plot( 129 | tf.squeeze(tests), 130 | tf.squeeze(pred_pso_gd2), 131 | "--", 132 | label="PSO - GD (lr= 1e-4)", 133 | ) 134 | plt.text(0.7, -1.2, "it: " + str(i * n_iter), fontsize=14) 135 | plt.text(-0.2, -1.4, "ANN layers: " + str(layer_sizes), fontsize=14) 136 | plt.legend() 137 | 138 | 139 | fig = plt.figure(figsize=(8, 8), dpi=100) 140 | anim = animation.FuncAnimation(fig, snapshot, frames=60) 141 | anim.save("swarm_training_PSOGD.gif", fps=6) 142 | -------------------------------------------------------------------------------- /src/swarm/optimizers/gd.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import tensorflow_probability as tfp 3 | from numpy import linspace 4 | from swarm import utils 5 | 6 | 7 | class gd: 8 | def __init__( 9 | self, 10 | loss_op, 11 | layer_sizes, 12 | n_iter=2000, 13 | pop_size=30, 14 | alpha=0.001, 15 | beta=0.9, 16 | x_min=-1, 17 | x_max=1, 18 | initialization_method="xavier", 19 | beta_1=0.99, 20 | beta_2=0.999, 21 | epsilon=1e-7, 22 | verbose=False, 23 | ): 24 | self.loss_op = loss_op 25 | self.layer_sizes = layer_sizes 26 | self.pop_size = pop_size 27 | self.n_iter = n_iter 28 | self.alpha = alpha 29 | self.beta = beta 30 | self.x_min = x_min 31 | self.x_max = x_max 32 | self.verbose = verbose 33 | self.initialization_method = initialization_method 34 | self.dim = utils.dimensions(layer_sizes) 35 | self.x = self.build_swarm() 36 | self.loss_history = [] 37 | self.beta_1 = beta_1 38 | self.beta_2 = beta_2 39 | self.epsilon = epsilon 40 | self.m1 = tf.zeros([self.pop_size, self.dim]) 41 | self.m2 = tf.zeros([self.pop_size, self.dim]) 42 | self.alpha = alpha * tf.math.sqrt(1 - self.beta_2) / (1 - self.beta_1) 43 | self.name = "ADAM" 44 | self.verbose_milestone = linspace(0, n_iter, 11).astype(int) 45 | 46 | def build_swarm(self): 47 | """Creates the swarm following the selected initialization method. 48 | 49 | Args: 50 | initialization_method (str): Chooses how to initialize the Neural Net weights. Allowed to be one of "uniform", "xavier", or "log_logistic". Defaults to None, where it uses uniform initialization. 51 | 52 | Returns: 53 | tf.Tensor: The PSO swarm population. Each particle represents a neural network. 54 | """ 55 | return utils.build_NN( 56 | self.pop_size, self.layer_sizes, self.initialization_method 57 | ) 58 | 59 | def individual_fn(self, particle): 60 | w, b = utils.decode(particle, self.layer_sizes) 61 | loss, grad = self.loss_op(w, b) 62 | return loss, utils.flat_grad(grad) 63 | 64 | @tf.function 65 | def fitness_fn(self, x): 66 | f_x, grads = tf.vectorized_map(self.individual_fn, x) 67 | return f_x[:, None], grads 68 | 69 | def adam_update(self): 70 | self.m1 = self.beta_1 * self.m1 + (1 - self.beta_1) * self.grads 71 | self.m2 = self.beta_2 * self.m2 + (1 - self.beta_2) * tf.math.square( 72 | self.grads 73 | ) 74 | return self.alpha * self.m1 / tf.math.sqrt(self.m2) + self.epsilon 75 | 76 | def first_momentum_update(self): 77 | self.m1 = self.beta * self.m1 + (1 - self.beta) * self.grads 78 | return self.alpha * self.m1 79 | 80 | def step(self): 81 | self.loss, self.grads = self.fitness_fn(self.x) 82 | self.x = self.x - self.adam_update() 83 | 84 | def train(self): 85 | for i in range(self.n_iter): 86 | self.step() 87 | self.loss_history.append(tf.reduce_mean(self.loss).numpy()) 88 | if self.verbose and i in self.verbose_milestone: 89 | utils.progress( 90 | (i / self.n_iter) * 100, 91 | metric="loss", 92 | metricValue=self.loss_history[-1], 93 | ) 94 | if self.verbose: 95 | utils.progress(100) 96 | print() 97 | 98 | def get_best(self): 99 | return utils.decode( 100 | self.x[tf.math.argmin(input=self.loss).numpy()[0]], 101 | self.layer_sizes, 102 | ) 103 | 104 | def get_swarm(self): 105 | return self.x 106 | -------------------------------------------------------------------------------- /examples/PINN/poisson/poisson.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import matplotlib as mpl 3 | import matplotlib.pyplot as plt 4 | import time 5 | import numpy as np 6 | import math 7 | import swarm.optimizers as optimizers 8 | from swarm.utils import multilayer_perceptron, decode, replacenan 9 | 10 | np.random.seed(123456) 11 | tf.random.set_seed(123456) 12 | 13 | uxn = 256 + 2 14 | xlo = 0 15 | xhi = 1 16 | 17 | ux = tf.reshape( 18 | tf.Variable(np.linspace(xlo, xhi, uxn), dtype="float32"), [uxn, 1] 19 | ) 20 | 21 | u = tf.sin(2 * math.pi * ux) 22 | 23 | x0 = tf.reshape(tf.convert_to_tensor(xlo, dtype=tf.float32), [1, 1]) 24 | x1 = tf.reshape(tf.convert_to_tensor(xhi, dtype=tf.float32), [1, 1]) 25 | u0 = tf.reshape(tf.convert_to_tensor(u[0], dtype=tf.float32), [1, 1]) 26 | u1 = tf.reshape(tf.convert_to_tensor(u[-1], dtype=tf.float32), [1, 1]) 27 | 28 | layer_sizes = [1] + 3 * [10] + [1] 29 | pop_size = 100 30 | n_iter = 30 31 | stepInd = 0.01 32 | stepVol = 0.01 33 | w_scale = 100 34 | 35 | 36 | def objective(x): 37 | return tf.sin(2 * math.pi * x) 38 | 39 | 40 | def u2(x): 41 | return tf.cast( 42 | -4 * np.pi * np.pi * tf.sin(2 * np.pi * x), dtype=tf.float32 43 | ) 44 | 45 | 46 | @tf.function 47 | def r(w, b): 48 | q = multilayer_perceptron(w, b, ux) 49 | q_x = tf.gradients(q, ux)[0] 50 | q_xx = tf.gradients(q_x, ux)[0] 51 | return tf.subtract(q_xx, u2(ux)) 52 | 53 | 54 | @tf.function 55 | def loss(w, b): 56 | pred_0 = multilayer_perceptron(w, b, x0) 57 | pred_1 = multilayer_perceptron(w, b, x1) 58 | pred_r = r(w, b) 59 | 60 | mse_0 = 100 * tf.pow(u0 - pred_0, 2) 61 | mse_1 = 100 * tf.pow(u1 - pred_1, 2) 62 | mse_r = tf.pow(pred_r, 2) 63 | 64 | return tf.reduce_mean(mse_0 + mse_1 + mse_r) 65 | 66 | 67 | def loss_grad(): 68 | def _loss(w, b): 69 | with tf.GradientTape(persistent=True) as tape: 70 | tape.watch(w) 71 | tape.watch(b) 72 | loss_value = loss(w, b) 73 | trainable_variables = w + b 74 | grads = tape.gradient(loss_value, trainable_variables) 75 | return loss_value, grads 76 | 77 | return _loss 78 | 79 | 80 | test_size = 100 81 | s = np.linspace(xlo, xhi, test_size) 82 | X = tf.reshape(tf.Variable(s, dtype="float32"), [test_size, 1]) 83 | 84 | opt = optimizers.get["pso_adam"]( 85 | loss_grad(), 86 | layer_sizes, 87 | n_iter, 88 | pop_size, 89 | 0.9, 90 | 0.08, 91 | 0.05, 92 | verbose=True, 93 | gd_alpha=1e-2, 94 | ) 95 | 96 | 97 | def nn_wrapper(particle): 98 | w, b = decode(particle, layer_sizes) 99 | return multilayer_perceptron(w, b, X) 100 | 101 | 102 | steps = 3 103 | mpl.style.use("seaborn") 104 | fig = plt.figure(figsize=(15, 7), dpi=300) 105 | plt.subplots_adjust(hspace=0.5) 106 | ax = [] 107 | 108 | for i in range(2): 109 | temp = [] 110 | for j in range(steps): 111 | temp.append(plt.subplot2grid((3, 3), (i, j))) 112 | ax.append(temp) 113 | ax.append(plt.subplot2grid((3, 3), (2, 0), colspan=3)) 114 | 115 | 116 | it_counter = n_iter 117 | for j in range(steps): 118 | start = time.time() 119 | opt.train() 120 | end = time.time() 121 | 122 | swarm = opt.get_swarm() 123 | preds = tf.reshape( 124 | tf.vectorized_map(nn_wrapper, swarm), [pop_size, test_size] 125 | ) 126 | 127 | mean = tf.reduce_mean(replacenan(preds), axis=0) 128 | variance = tf.math.reduce_variance(replacenan(preds), axis=0) 129 | 130 | nn_w, nn_b = opt.get_best() 131 | pred = multilayer_perceptron(nn_w, nn_b, X) 132 | X_ = tf.squeeze(X) 133 | y = objective(X) 134 | 135 | for k in range(preds.shape[0]): 136 | ax[0][j].plot(X_, preds[k], linewidth=0.2) 137 | ax[1][j].fill_between( 138 | X_, 139 | mean - variance, 140 | mean + variance, 141 | color="gray", 142 | alpha=0.5, 143 | ) 144 | ax[1][j].plot(X_, y, label="Original Function") 145 | ax[1][j].plot(X_, pred, "--", label="Best") 146 | ax[1][j].plot(X_, mean, "--", label="Mean") 147 | ax[0][j].set_title(str(it_counter) + " iterations", fontsize="xx-large") 148 | ax[1][j].legend() 149 | 150 | it_counter = it_counter + n_iter 151 | 152 | 153 | ax[2].plot(opt.loss_history) 154 | ax[2].set_title("Loss Error (mean)", fontsize="xx-large") 155 | 156 | plt.savefig("poisson_train.png", bbox_inches="tight") 157 | -------------------------------------------------------------------------------- /examples/PINN/poisson/poisson_ani.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import matplotlib.pyplot as plt 3 | import matplotlib.animation as animation 4 | import numpy as np 5 | import math 6 | 7 | from swarm.optimizers.pso import pso 8 | 9 | from swarm.utils import multilayer_perceptron 10 | 11 | 12 | np.random.seed(1234) 13 | tf.random.set_seed(1234) 14 | 15 | uxn = 256 + 2 16 | xlo = 0 17 | xhi = 1 18 | 19 | ux = tf.reshape(tf.Variable(np.linspace(xlo, xhi, uxn), dtype="float32"), [uxn, 1]) 20 | 21 | u = tf.sin(2 * math.pi * ux) 22 | 23 | x0 = tf.reshape(tf.convert_to_tensor(xlo, dtype=tf.float32), [1, 1]) 24 | x1 = tf.reshape(tf.convert_to_tensor(xhi, dtype=tf.float32), [1, 1]) 25 | u0 = tf.reshape(tf.convert_to_tensor(u[0], dtype=tf.float32), [1, 1]) 26 | u1 = tf.reshape(tf.convert_to_tensor(u[-1], dtype=tf.float32), [1, 1]) 27 | 28 | layer_sizes = [1] + 3 * [20] + [1] 29 | pop_size = 100 30 | n_iter = 10 31 | stepInd = 0.01 32 | stepVol = 0.01 33 | w_scale = 100 34 | x_min = -1 35 | x_max = 1 36 | sample_size = 10 37 | 38 | 39 | noise = 0.0 40 | 41 | 42 | def objective(x, noise=0): 43 | return (tf.sin(2 * math.pi * x)) + tf.random.normal(x.shape, 0, noise, tf.float32) 44 | 45 | 46 | def u2(x): 47 | return tf.cast(-4 * np.pi * np.pi * tf.sin(2 * np.pi * x), dtype=tf.float32) 48 | 49 | 50 | @tf.function 51 | def r(w, b): 52 | q = multilayer_perceptron(w, b, ux) 53 | q_x = tf.gradients(q, ux)[0] 54 | q_xx = tf.gradients(q_x, ux)[0] 55 | return tf.subtract(q_xx, u2(ux)) 56 | 57 | 58 | @tf.function 59 | def loss(w, b): 60 | pred_0 = multilayer_perceptron(w, b, x0) 61 | pred_1 = multilayer_perceptron(w, b, x1) 62 | pred_r = r(w, b) 63 | 64 | mse_0 = 100 * tf.pow(u0 - pred_0, 2) 65 | mse_1 = 100 * tf.pow(u1 - pred_1, 2) 66 | mse_r = tf.pow(pred_r, 2) 67 | 68 | return tf.reduce_mean(mse_0 + mse_1 + mse_r) 69 | 70 | 71 | def loss_grad(): 72 | def _loss(w, b): 73 | with tf.GradientTape(persistent=True) as tape: 74 | tape.watch(w) 75 | tape.watch(b) 76 | loss_value = loss(w, b) 77 | 78 | trainable_variables = w + b 79 | grads = tape.gradient(loss_value, trainable_variables) 80 | return loss_value, grads 81 | 82 | return _loss 83 | 84 | 85 | opt_pso = pso( 86 | loss_grad(), 87 | layer_sizes, 88 | n_iter, 89 | pop_size, 90 | 0.9, 91 | 0.8, 92 | 0.5, 93 | x_min=x_min, 94 | x_max=x_max, 95 | gd_alpha=0, 96 | ) 97 | 98 | opt_pso_gd_1 = pso( 99 | loss_grad(), 100 | layer_sizes, 101 | n_iter, 102 | pop_size, 103 | 0.9, 104 | 0.8, 105 | 0.5, 106 | x_min=x_min, 107 | x_max=x_max, 108 | gd_alpha=1e-3, 109 | ) 110 | 111 | opt_pso_gd_2 = pso( 112 | loss_grad(), 113 | layer_sizes, 114 | n_iter, 115 | pop_size, 116 | 0.9, 117 | 0.8, 118 | 0.5, 119 | x_min=x_min, 120 | x_max=x_max, 121 | gd_alpha=1e-4, 122 | ) 123 | 124 | opt_pso_gd_3 = pso( 125 | loss_grad(), 126 | layer_sizes, 127 | n_iter, 128 | pop_size, 129 | 0.9, 130 | 0.8, 131 | 0.5, 132 | x_min=x_min, 133 | x_max=x_max, 134 | gd_alpha=1e-5, 135 | ) 136 | 137 | 138 | s = np.linspace(xlo, xhi, 100) 139 | tests = tf.reshape(tf.Variable(s, dtype="float32"), [100, 1]) 140 | results = objective(tests) 141 | 142 | 143 | def snapshot(i): 144 | 145 | opt_pso.train() 146 | nn_w_pso, nn_b_pso = opt_pso.get_best() 147 | pred_pso = multilayer_perceptron(nn_w_pso, nn_b_pso, tests) 148 | 149 | opt_pso_gd_1.train() 150 | nn_w_pso_1, nn_b_pso_1 = opt_pso_gd_1.get_best() 151 | pred_pso_1 = multilayer_perceptron(nn_w_pso_1, nn_b_pso_1, tests) 152 | 153 | opt_pso_gd_2.train() 154 | nn_w_pso_2, nn_b_pso_2 = opt_pso_gd_2.get_best() 155 | pred_pso_2 = multilayer_perceptron(nn_w_pso_2, nn_b_pso_2, tests) 156 | 157 | opt_pso_gd_3.train() 158 | nn_w_pso_3, nn_b_pso_3 = opt_pso_gd_3.get_best() 159 | pred_pso_3 = multilayer_perceptron(nn_w_pso_3, nn_b_pso_3, tests) 160 | 161 | plt.clf() 162 | plt.xlim([0, 1]) 163 | plt.ylim([-1.5, 1.5]) 164 | 165 | plt.plot(tf.squeeze(tests), tf.squeeze(results), label="$f(x) = \sin (2 \pi x)$") 166 | plt.plot(tf.squeeze(tests), tf.squeeze(pred_pso), "--", label="PSO ") 167 | plt.plot(tf.squeeze(tests), tf.squeeze(pred_pso_1), "--", label="PSO GD (lr: 1e-3)") 168 | plt.plot( 169 | tf.squeeze(tests), tf.squeeze(pred_pso_2), "--", label="PSO - GD (lr: 1e-4)" 170 | ) 171 | plt.plot( 172 | tf.squeeze(tests), 173 | tf.squeeze(pred_pso_3), 174 | "--", 175 | label="PSO - GD (lr: 1e-5)", 176 | ) 177 | plt.text(0.7, -1.2, "it: " + str(i * n_iter), fontsize=14) 178 | plt.legend() 179 | 180 | 181 | fig = plt.figure(figsize=(8, 8), dpi=100) 182 | anim = animation.FuncAnimation(fig, snapshot, frames=40) 183 | anim.save("poisson_PSO.gif", fps=2) 184 | -------------------------------------------------------------------------------- /src/swarm/optimizers/fss.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from swarm import utils 3 | 4 | import sys 5 | 6 | 7 | class fss: 8 | def __init__( 9 | self, 10 | loss_op, 11 | layer_sizes, 12 | iter=2000, 13 | pop_size=30, 14 | w_scale=100, 15 | stepInd=0.01, 16 | stepVol=0.01, 17 | x_min=-1, 18 | x_max=1, 19 | xavier_init=False, 20 | verbose=False, 21 | newOP=False, 22 | ): 23 | self.loss_op = loss_op 24 | self.layer_sizes = layer_sizes 25 | self.iter = iter 26 | self.pop_size = pop_size 27 | self.stepInd = stepInd 28 | self.stepIndDecay = self.stepInd / self.iter 29 | self.stepVol = stepVol 30 | self.stepVolDecay = self.stepVol / self.iter 31 | self.w_scale = w_scale 32 | self.x_min = x_min 33 | self.x_max = x_max 34 | self.xavier_init = xavier_init 35 | self.dim = utils.dimensions(layer_sizes) 36 | self.X, self.w = self.make_pop() 37 | self.f_X = self.f(self.X) 38 | self.verbose = verbose 39 | self.newOP = newOP 40 | self.loss_history = [] 41 | 42 | def individual_fn(self, particle): 43 | w, b = utils.decode(particle, self.layer_sizes) 44 | loss, _ = self.loss_op(w, b) 45 | return -loss 46 | 47 | @tf.function 48 | def f(self, x): 49 | f_x = tf.vectorized_map(self.individual_fn, x) 50 | return f_x[:, None] 51 | 52 | def individual(self): 53 | step = ( 54 | tf.random.uniform([self.pop_size, self.dim], self.x_min, self.x_max) 55 | * self.stepInd 56 | ) 57 | x1 = tf.add(self.X, step) 58 | f_x1 = self.f(x1) 59 | x1 = tf.where(f_x1 > self.f_X, x1, self.X) 60 | f_x1 = tf.where(f_x1 > self.f_X, f_x1, self.f_X) 61 | step = tf.where(f_x1 > self.f_X, step, tf.zeros([self.pop_size, self.dim])) 62 | return x1, step, f_x1 63 | 64 | def instictive(self, x1, step, f_x1): 65 | if self.newOP: 66 | self.__instictive__(x1, step) 67 | else: 68 | self._instictive_(x1, step, f_x1) 69 | 70 | def _instictive_(self, x1, step, f_x1): 71 | sum = tf.add(f_x1, -self.f_X) 72 | self.X = tf.add( 73 | x1, 74 | utils.replacenan( 75 | tf.divide(tf.reduce_sum(step * sum, axis=0), tf.reduce_sum(sum)) 76 | ), 77 | ) 78 | self.f_X = self.f(self.X) 79 | self.loss_history.append(-tf.reduce_min(self.f_X).numpy()) 80 | 81 | def __instictive__(self, x1, step): 82 | self.X = tf.add(x1, step) 83 | self.f_X = self.f(self.X) 84 | self.loss_history.append(-tf.reduce_min(self.f_X).numpy()) 85 | 86 | def bari(self): 87 | den = tf.reduce_sum(self.w[:, None] * self.X, 0) 88 | return den / tf.reduce_sum(self.w) 89 | 90 | def feed(self, f_x1): 91 | df = self.f_X - f_x1 92 | df_mean = df / tf.reduce_max(df) 93 | w1 = tf.add(self.w, tf.squeeze(df_mean)) 94 | return tf.clip_by_value(w1, 0, self.w_scale) 95 | 96 | def volitive(self, w1): 97 | rand = tf.scalar_mul(self.stepVol, tf.random.uniform([self.pop_size, 1], 0, 1)) 98 | bari_vector = utils.replacenan(tf.add(self.X, -self.bari())) 99 | step = tf.multiply(rand, bari_vector) 100 | x_contract = tf.add(self.X, -step) 101 | x_dilate = tf.add(self.X, step) 102 | self.X = tf.where(w1[:, None] > self.w[:, None], x_contract, x_dilate) 103 | self.w = w1 104 | 105 | def make_pop(self): 106 | if self.xavier_init: 107 | X = utils.make_xavier_NN(self.pop_size,self.layer_sizes) 108 | else: 109 | X = tf.Variable( 110 | tf.random.uniform([self.pop_size, self.dim], self.x_min, self.x_max) 111 | ) 112 | w = tf.Variable([self.w_scale / 2] * self.pop_size) 113 | return X, w 114 | 115 | def update_steps(self): 116 | self.stepInd = self.stepInd - self.stepIndDecay 117 | self.stepVol = self.stepVol - self.stepVolDecay 118 | 119 | def train(self): 120 | for i in range(self.iter): 121 | x1, step, f_x1 = self.individual() 122 | w1 = self.feed(f_x1) 123 | self.instictive(x1, step, f_x1) 124 | self.volitive(w1) 125 | self.update_steps() 126 | if self.verbose and i % (self.iter / 10) == 0: 127 | utils.progress((i / self.iter) * 100) 128 | if self.verbose: 129 | utils.progress(100) 130 | 131 | def get_best(self): 132 | return utils.decode( 133 | tf.unstack(self.X)[tf.math.argmax(tf.reshape(self.f_X, [self.pop_size]))], 134 | self.layer_sizes, 135 | ) 136 | 137 | def get_swarm(self): 138 | return self.X 139 | -------------------------------------------------------------------------------- /examples/PINN/advection/advection.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | from pyDOE import lhs 4 | import matplotlib as mpl 5 | import matplotlib.pyplot as plt 6 | 7 | from swarm.optimizers.pso import pso 8 | from swarm.utils import multilayer_perceptron 9 | 10 | np.random.seed(123456) 11 | tf.random.set_seed(123456) 12 | 13 | # parameters of the simulation 14 | ql = 4 # concentration at the left 15 | qm1 = 0.4 # intermediate concentration 16 | qm2 = 4 # intermediate concentration 17 | qm3 = 1.4 # intermediate concentration 18 | qm4 = 3 # intermediate concentration 19 | qm5 = 0.4 # intermediate concentration 20 | qm6 = 4 # intermediate concentration 21 | qr = 0.4 # concentration at the right 22 | x0 = 1 # central discontinuity location 23 | x01 = 0.25 # left discontinuity location 24 | x02 = 0.5 # right discontinuity location 25 | x03 = 0.75 # left discontinuity location 26 | x04 = 1.25 # right discontinuity location 27 | x05 = 1.5 # left discontinuity location 28 | x06 = 1.75 # right discontinuity location 29 | L = 2 # channel length 30 | u = 1 # advection speed 31 | 32 | # nu = 0.01 # viscosity parameter 33 | 34 | # define grid for solution 35 | utn = 256 36 | uxn = 1024 37 | xlo = 0 38 | xhi = L 39 | ux = np.linspace(xlo, xhi, uxn) 40 | tlo = 0.0 41 | thi = 0.2 42 | ut = np.linspace(tlo, thi, utn) 43 | 44 | # analytical solution basic step 45 | q = np.zeros([uxn, utn]) 46 | for utj in range(utn): 47 | j0 = int(uxn * (x0 + u * ut[utj]) / L) 48 | for uxj in range(0, j0 + 1): 49 | q[uxj, utj] = ql 50 | for uxj in range(j0 + 1, uxn): 51 | q[uxj, utj] = qr 52 | 53 | layer_sizes = [2] + 5 * [10] + [1] 54 | pop_size = 100 55 | n_iter = 100 56 | 57 | 58 | # collocation points 59 | Ncl = 5000 60 | X = lhs(2, Ncl) 61 | xcl = tf.expand_dims( 62 | tf.convert_to_tensor(xlo + (xhi - xlo) * X[:, 0], dtype=tf.float32), -1 63 | ) 64 | tcl = tf.expand_dims( 65 | tf.convert_to_tensor(tlo + (thi - tlo) * X[:, 1], dtype=tf.float32), -1 66 | ) 67 | 68 | # initial condition points 69 | x0 = tf.expand_dims(tf.convert_to_tensor(ux, dtype=tf.float32), -1) 70 | t0 = tf.zeros(tf.shape(x0), dtype=tf.float32) 71 | q0 = tf.expand_dims(tf.convert_to_tensor(q[:, 0], dtype=tf.float32), -1) 72 | 73 | # Dirichlet boundary condition points 74 | xlb = tf.expand_dims(xlo * tf.ones(tf.shape(ut), dtype=tf.float32), -1) 75 | tlb = tf.expand_dims(tf.convert_to_tensor(ut, dtype=tf.float32), -1) 76 | qlb = ql * tf.ones(tf.shape(tlb), dtype=tf.float32) 77 | 78 | xub = tf.expand_dims(xhi * tf.ones(tf.shape(ut), dtype=tf.float32), -1) 79 | tub = tf.expand_dims(tf.convert_to_tensor(ut, dtype=tf.float32), -1) 80 | qub = qr * tf.ones(tf.shape(tub), dtype=tf.float32) 81 | 82 | 83 | def f(w, b, x, t): 84 | q = multilayer_perceptron(w, b, tf.concat([x, t], 1)) # compute q 85 | 86 | q_t = tf.gradients(q, t)[0] 87 | q_x = tf.gradients(q, x)[0] 88 | 89 | fr = q_t + u * q_x 90 | 91 | return fr 92 | 93 | 94 | @tf.function 95 | def loss(w, b): 96 | 97 | pred_0 = multilayer_perceptron(w, b, tf.concat([x0, t0], 1)) 98 | pred_lb = multilayer_perceptron(w, b, tf.concat([xlb, tlb], 1)) 99 | pred_ub = multilayer_perceptron(w, b, tf.concat([xub, tub], 1)) 100 | fr = f(w, b, xcl, tcl) 101 | 102 | # IC loss 103 | mse_0 = tf.reduce_mean(tf.pow(q0 - pred_0, 2)) 104 | 105 | # Dirichlet boundary loss 106 | mse_lb = tf.reduce_mean(tf.pow(qlb - pred_lb, 2)) 107 | mse_ub = tf.reduce_mean(tf.pow(qub - pred_ub, 2)) 108 | 109 | # Residual loss 110 | mse_f = tf.reduce_mean(tf.pow(fr, 2)) 111 | 112 | return mse_0 + mse_f + mse_lb + mse_ub 113 | 114 | 115 | def loss_grad(): 116 | def _loss(w, b): 117 | with tf.GradientTape(persistent=True) as tape: 118 | tape.watch(w) 119 | tape.watch(b) 120 | loss_value = loss(w, b) 121 | trainable_variables = w + b 122 | grads = tape.gradient(loss_value, trainable_variables) 123 | return loss_value, grads 124 | 125 | return _loss 126 | 127 | 128 | opt = pso( 129 | loss_grad(), 130 | layer_sizes, 131 | n_iter, 132 | pop_size, 133 | 0.9, 134 | 0.08, 135 | 0.05, 136 | verbose=True, 137 | gd_alpha=1e-2, 138 | ) 139 | 140 | opt.train() 141 | 142 | X, T = np.meshgrid(ux, ut) 143 | X_flat = np.hstack((X.flatten()[:, None], T.flatten()[:, None])) 144 | nn_w, nn_b = opt.get_best() 145 | u_pso = multilayer_perceptron(nn_w, nn_b, X_flat.astype(np.float32)) 146 | 147 | q_flat = q.T.flatten() 148 | err_pso = np.linalg.norm(q_flat - tf.squeeze(u_pso), 2) / np.linalg.norm(q_flat, 2) 149 | print("L2 error PSO-GD: %.2e" % (err_pso)) 150 | 151 | mpl.style.use("seaborn") 152 | fig = plt.figure(figsize=(8, 8), dpi=300) 153 | plt.subplots_adjust(hspace=0.3) 154 | ax = fig.subplots(3) 155 | m = [0, int(utn / 2), utn - 1] 156 | 157 | for i in [0, 1, 2]: 158 | ax[i].plot(q[:, m[i]], "-", linewidth=3, label="Analytical") 159 | ax[i].plot( 160 | u_pso[m[i] * uxn : (m[i] + 1) * uxn], "--", linewidth=3, label="PSO-PINN" 161 | ) 162 | ax[i].set_title(r"t = {:.2f}/$\pi$".format(np.pi * m[i] * thi / (utn - 1))) 163 | ax[i].set_ylabel("$u(t,x)$") 164 | ax[i].legend() 165 | ax[2].set_xlabel("$x$", fontsize="xx-large") 166 | 167 | plt.savefig("advection.png", bbox_inches="tight") 168 | -------------------------------------------------------------------------------- /examples/PINN/diffusion/diffusion.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | import argparse 4 | import time 5 | 6 | from pyDOE import lhs 7 | 8 | import swarm.optimizers as optimizers 9 | 10 | 11 | import matplotlib.pyplot as plt 12 | from swarm.utils import multilayer_perceptron, decode 13 | import matplotlib.animation as animation 14 | 15 | 16 | from scipy.interpolate import griddata 17 | 18 | 19 | # define grid for solution 20 | utn = 256 21 | uxn = 1024 22 | xlo = 0 23 | xhi = 1 24 | ux = np.linspace(xlo, xhi, uxn) 25 | tlo = 0 26 | thi = 1 27 | ut = np.linspace(tlo, thi, utn) 28 | 29 | 30 | def difusion_eq_exact_solution(x, t): 31 | return ( 32 | 1 33 | + t 34 | + np.exp(-4 * (np.pi**2) * t) * np.cos(2 * np.pi * x) 35 | + x * np.sin(t) 36 | ) 37 | 38 | 39 | u_quad = [] 40 | for utj in ut: 41 | ux_t = np.linspace(xlo, xhi, uxn) 42 | u_quad.append(difusion_eq_exact_solution(ux_t, utj)) 43 | 44 | u_quad = np.array(u_quad).T 45 | 46 | # collocation points 47 | Ncl = 512 48 | X = lhs(2, Ncl) 49 | xcl = tf.expand_dims( 50 | tf.convert_to_tensor(xlo + (xhi - xlo) * X[:, 0], dtype=tf.float32), -1 51 | ) 52 | tcl = tf.expand_dims( 53 | tf.convert_to_tensor(tlo + (thi - tlo) * X[:, 1], dtype=tf.float32), -1 54 | ) 55 | 56 | # initial condition points 57 | x0 = tf.expand_dims(tf.convert_to_tensor(ux, dtype=tf.float32), -1) 58 | t0 = tf.zeros(tf.shape(x0), dtype=tf.float32) 59 | q0 = tf.convert_to_tensor(1 + np.cos(2 * np.pi * x0), dtype=tf.float32) 60 | 61 | # Dirichlet boundary condition points 62 | xlb = tf.expand_dims(xlo * tf.ones(tf.shape(ut), dtype=tf.float32), -1) 63 | tlb = tf.expand_dims(tf.convert_to_tensor(ut, dtype=tf.float32), -1) 64 | 65 | xub = tf.expand_dims(xhi * tf.ones(tf.shape(ut), dtype=tf.float32), -1) 66 | tub = tf.expand_dims(tf.convert_to_tensor(ut, dtype=tf.float32), -1) 67 | 68 | qb = tf.convert_to_tensor(np.sin(tlb), dtype=tf.float32) 69 | 70 | # residual 71 | def get_residual(x, t): 72 | return 1 + x * np.cos(t) 73 | 74 | 75 | res = [] 76 | for tcli in tcl: 77 | res.append(get_residual(xcl, tcli)) 78 | 79 | res = np.asarray(res) 80 | 81 | 82 | @tf.function 83 | def f(w, b, x, t): 84 | u = multilayer_perceptron(w, b, tf.concat([x, t], 1)) 85 | u_x = tf.gradients(u, x)[0] 86 | u_xx = tf.gradients(u_x, x)[0] 87 | u_t = tf.gradients(u, t)[0] 88 | f_u = u_t - u_xx - res 89 | return f_u 90 | 91 | 92 | @tf.function 93 | def fx(w, b, x, t): 94 | u = multilayer_perceptron(w, b, tf.concat([x, t], 1)) 95 | u_x = tf.gradients(u, x)[0] 96 | return u_x 97 | 98 | 99 | def loss(w, b): 100 | pred_0 = multilayer_perceptron(w, b, tf.concat([x0, t0], 1)) 101 | pred_x0 = fx(w, b, xlb, tlb) 102 | pred_x1 = fx(w, b, xub, tub) 103 | fr = f(w, b, xcl, tcl) 104 | 105 | # IC loss 106 | mse_0 = tf.reduce_mean(tf.pow(q0 - pred_0, 2)) 107 | 108 | # Dirichlet boundary loss 109 | mse_lb = tf.reduce_mean(tf.pow(qb - pred_x0, 2)) 110 | mse_ub = tf.reduce_mean(tf.pow(qb - pred_x1, 2)) 111 | 112 | # Residual loss 113 | mse_f = tf.reduce_mean(tf.pow(fr, 2)) 114 | 115 | return mse_0 + mse_f + mse_lb + mse_ub 116 | 117 | 118 | def loss_grad(): 119 | def _loss(w, b): 120 | with tf.GradientTape(persistent=True) as tape: 121 | tape.watch(w) 122 | tape.watch(b) 123 | loss_value = loss(w, b) 124 | trainable_variables = w + b 125 | grads = tape.gradient(loss_value, trainable_variables) 126 | return loss_value, grads 127 | 128 | return _loss 129 | 130 | 131 | optmizer = "pso_adam" 132 | hidden_layers = 5 133 | neurons_per_layer = 8 134 | layer_sizes = [2] + hidden_layers * [neurons_per_layer] + [1] 135 | n_iter = 4000 136 | 137 | 138 | def run_swarm(swarm, X): 139 | new_swarm = [] 140 | for particle in swarm: 141 | w, b = decode(particle, layer_sizes) 142 | new_swarm.append(multilayer_perceptron(w, b, X)) 143 | return new_swarm 144 | 145 | 146 | X, T = np.meshgrid(ux, ut) 147 | X_flat = tf.convert_to_tensor( 148 | np.hstack((X.flatten()[:, None], T.flatten()[:, None])), 149 | dtype=tf.float32, 150 | ) 151 | q_flat = u_quad.T.flatten() 152 | 153 | opt = optimizers.get[optmizer](loss_grad(), layer_sizes, n_iter) 154 | 155 | start = time.time() 156 | opt.train() 157 | end = time.time() 158 | swarm = opt.get_swarm() 159 | preds = run_swarm(swarm, X_flat) 160 | mean = tf.squeeze(tf.reduce_mean(preds, axis=0)) 161 | var = tf.squeeze(tf.math.reduce_std(preds, axis=0)) 162 | err = np.linalg.norm(q_flat - mean, 2) / np.linalg.norm(q_flat, 2) 163 | 164 | time_steps = utn # total number of time steps in animation 165 | fps = 15 # frames/second of animation 166 | 167 | 168 | def snapshot(i): 169 | l_ind = i * uxn 170 | u_ind = (i + 1) * uxn 171 | plt.clf() 172 | plt.ylim([0, 3]) 173 | plt.plot(ux, q_flat[l_ind:u_ind], "b-", linewidth=3, label="Quadrature") 174 | locs, labels = plt.xticks() 175 | plt.plot( 176 | ux, 177 | mean[l_ind:u_ind], 178 | "r--", 179 | linewidth=3, 180 | label=opt.name, 181 | ) 182 | plt.fill_between( 183 | ux, 184 | mean[l_ind:u_ind] - 3 * var[l_ind:u_ind], 185 | mean[l_ind:u_ind] + 3 * var[l_ind:u_ind], 186 | color="gray", 187 | alpha=0.2, 188 | ) 189 | plt.xlabel("$x$", fontsize="xx-large") 190 | plt.ylabel("$u(t,x)$", fontsize="xx-large") 191 | plt.grid() 192 | plt.legend(fontsize="x-large") 193 | 194 | 195 | fig = plt.figure(figsize=(8, 8), dpi=150) 196 | anim = animation.FuncAnimation(fig, snapshot, frames=time_steps) 197 | anim.save("Diffusion_Demo.gif", fps=fps) 198 | -------------------------------------------------------------------------------- /examples/PINN/burgers/burgers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy.polynomial.hermite import hermgauss 3 | import tensorflow as tf 4 | import time 5 | from pyDOE import lhs 6 | from swarm.optimizers.pso_adam import pso 7 | from swarm.utils import multilayer_perceptron, decode 8 | import matplotlib.pyplot as plt 9 | import matplotlib.animation as animation 10 | 11 | np.random.seed(12345) 12 | tf.random.set_seed(12345) 13 | 14 | 15 | """ 16 | Problem Definition and Quadrature Solution 17 | """ 18 | 19 | nu = 0.01 # / np.pi # viscosity parameter 20 | 21 | # define grid for quadrature solution 22 | utn = 100 23 | uxn = 256 24 | xlo = -1.0 25 | xhi = +1.0 26 | ux = np.linspace(xlo, xhi, uxn) 27 | tlo = 0.0 28 | thi = 1.0 29 | ut = np.linspace(tlo, thi, utn) 30 | 31 | qn = 32 # order of quadrature rule 32 | qx, qw = hermgauss(qn) 33 | 34 | # compute solution u(x,t) by quadrature of analytical formula: 35 | u_quad = np.zeros([uxn, utn]) 36 | for utj in range(utn): 37 | if ut[utj] == 0.0: 38 | for uxj in range(uxn): 39 | u_quad[uxj, utj] = -np.sin(np.pi * ux[uxj]) 40 | else: 41 | for uxj in range(uxn): 42 | top = 0.0 43 | bot = 0.0 44 | for qj in range(qn): 45 | c = 2.0 * np.sqrt(nu * ut[utj]) 46 | top = top - qw[qj] * c * np.sin( 47 | np.pi * (ux[uxj] - c * qx[qj]) 48 | ) * np.exp( 49 | -np.cos(np.pi * (ux[uxj] - c * qx[qj])) 50 | / (2.0 * np.pi * nu) 51 | ) 52 | bot = bot + qw[qj] * c * np.exp( 53 | -np.cos(np.pi * (ux[uxj] - c * qx[qj])) 54 | / (2.0 * np.pi * nu) 55 | ) 56 | u_quad[uxj, utj] = top / bot 57 | 58 | 59 | # Algorithm parameters 60 | layer_sizes = [2] + 5 * [15] + [1] 61 | pop_size = 10 62 | n_iter = 6000 63 | 64 | 65 | # collocation points 66 | def collocation_points(size): 67 | X = lhs(2, size) 68 | xcl = tf.expand_dims( 69 | tf.convert_to_tensor(xlo + (xhi - xlo) * X[:, 0], dtype=tf.float32), -1 70 | ) 71 | tcl = tf.expand_dims( 72 | tf.convert_to_tensor(tlo + (thi - tlo) * X[:, 1], dtype=tf.float32), -1 73 | ) 74 | return xcl, tcl 75 | 76 | 77 | # initial condition points 78 | x0 = tf.expand_dims(tf.convert_to_tensor(ux, dtype=tf.float32), -1) 79 | t0 = tf.zeros(tf.shape(x0), dtype=tf.float32) 80 | u0 = -tf.math.sin(np.pi * x0) 81 | 82 | # Dirichlet boundary condition points 83 | xlb = tf.expand_dims(xlo * tf.ones(tf.shape(ut), dtype=tf.float32), -1) 84 | tlb = tf.expand_dims(tf.convert_to_tensor(ut, dtype=tf.float32), -1) 85 | ulb = tf.expand_dims(tf.zeros(tf.shape(ut), dtype=tf.float32), -1) 86 | xub = tf.expand_dims(xhi * tf.ones(tf.shape(ut), dtype=tf.float32), -1) 87 | tub = tf.expand_dims(tf.convert_to_tensor(ut, dtype=tf.float32), -1) 88 | uub = tf.expand_dims(tf.zeros(tf.shape(ut), dtype=tf.float32), -1) 89 | 90 | 91 | @tf.function 92 | def f_model(w, b): 93 | x, t = collocation_points(500) 94 | u = multilayer_perceptron(w, b, tf.concat([x, t], 1)) 95 | u_x = tf.gradients(u, x) 96 | u_xx = tf.gradients(u_x, x) 97 | u_t = tf.gradients(u, t) 98 | f_u = u_t + u * u_x - [nu * element for element in u_xx] 99 | return f_u 100 | 101 | 102 | @tf.function 103 | def loss(w, b): 104 | u0_pred = multilayer_perceptron(w, b, tf.concat([x0, t0], 1)) 105 | ulb_pred = multilayer_perceptron(w, b, tf.concat([xlb, tlb], 1)) 106 | uub_pred = multilayer_perceptron(w, b, tf.concat([xub, tub], 1)) 107 | f_pred = f_model(w, b) 108 | 109 | # IC loss 110 | mse_0 = tf.reduce_mean(tf.pow(u0 - u0_pred, 2)) 111 | 112 | # Dirichlet BC loss 113 | mse_lb = tf.reduce_mean(tf.pow(ulb_pred - ulb, 2)) 114 | mse_ub = tf.reduce_mean(tf.pow(uub_pred - uub, 2)) 115 | 116 | # Residual loss 117 | mse_f = tf.reduce_mean(tf.pow(f_pred, 2)) 118 | 119 | return 1.5 * mse_0 + mse_f + mse_lb + mse_ub 120 | 121 | 122 | def loss_grad(): 123 | def _loss(w, b): 124 | with tf.GradientTape(persistent=True) as tape: 125 | tape.watch(w) 126 | tape.watch(b) 127 | loss_value = loss(w, b) 128 | trainable_variables = w + b 129 | grads = tape.gradient(loss_value, trainable_variables) 130 | return loss_value, grads 131 | 132 | return _loss 133 | 134 | 135 | def run_swarm(swarm, X_flat): 136 | new_swarm = [] 137 | for particle in swarm: 138 | w, b = decode(particle, layer_sizes) 139 | new_swarm.append( 140 | multilayer_perceptron(w, b, X_flat.astype(np.float32)) 141 | ) 142 | return tf.convert_to_tensor(new_swarm, dtype=tf.float32) 143 | 144 | 145 | opt = pso( 146 | loss_grad(), 147 | layer_sizes, 148 | n_iter, 149 | pop_size, 150 | 0.999, 151 | 8e-4, 152 | 5e-3, 153 | initialization_method="xavier", 154 | verbose=True, 155 | gd_alpha=1e-4, 156 | ) 157 | 158 | start = time.time() 159 | opt.train() 160 | end = time.time() 161 | print("\nTime elapsed: ", end - start) 162 | 163 | X, T = np.meshgrid(ux, ut) 164 | X_flat = np.hstack((X.flatten()[:, None], T.flatten()[:, None])) 165 | 166 | nn_w, nn_b = opt.get_best() 167 | best = multilayer_perceptron(nn_w, nn_b, X_flat.astype(np.float32)) 168 | 169 | swarm = opt.get_swarm() 170 | preds = run_swarm(swarm, X_flat.astype(np.float32)) 171 | mean = tf.squeeze(tf.reduce_mean(preds, axis=0)) 172 | var = tf.squeeze(tf.math.reduce_std(preds, axis=0)) 173 | 174 | u_PINN = multilayer_perceptron(nn_w, nn_b, X_flat.astype(np.float32)) 175 | 176 | """ 177 | Comparison and Animation 178 | """ 179 | 180 | u_quad_flat = u_quad.T.flatten()[:, None] 181 | error_u = np.linalg.norm(u_quad_flat - u_PINN, 2) / np.linalg.norm( 182 | u_quad_flat, 2 183 | ) 184 | print("nu = %.2f/pi" % (np.pi * nu)) 185 | print("L2 Error = %e" % (error_u)) 186 | print("Last Loss: ", opt.loss_history[-1]) 187 | 188 | time_steps = utn # total number of time steps in animation 189 | fps = 15 # frames/second of animation 190 | 191 | 192 | def snapshot(i): 193 | plt.clf() 194 | plt.plot(ux, u_quad[:, i], "b-", linewidth=3, label="Quadrature") 195 | plt.plot( 196 | ux, 197 | mean[i * uxn : (i + 1) * uxn], 198 | "r--", 199 | linewidth=3, 200 | label="PSO-PINN", 201 | ) 202 | plt.fill_between( 203 | ux, 204 | mean[i * uxn : (i + 1) * uxn] - var[i * uxn : (i + 1) * uxn], 205 | mean[i * uxn : (i + 1) * uxn] + var[i * uxn : (i + 1) * uxn], 206 | color="gray", 207 | alpha=0.5, 208 | ) 209 | plt.xlabel("$x$", fontsize="xx-large") 210 | plt.ylabel("$u(t,x)$", fontsize="xx-large") 211 | plt.grid() 212 | plt.xlim(-1.02, 1.02) 213 | plt.ylim(-1.02, 1.02) 214 | plt.legend(fontsize="x-large") 215 | 216 | 217 | fig = plt.figure(figsize=(8, 8), dpi=150) 218 | # Call the animator: 219 | anim = animation.FuncAnimation(fig, snapshot, frames=time_steps) 220 | # Save the animation as an mp4. This requires ffmpeg to be installed. 221 | anim.save("Burgers_Demo.gif", fps=fps) 222 | -------------------------------------------------------------------------------- /examples/PINN/allen-cahn/ac.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import time 5 | import pickle 6 | 7 | from swarm.optimizers.pso_adam import pso 8 | from swarm.utils import multilayer_perceptron, decode 9 | import matplotlib.pyplot as plt 10 | import matplotlib.animation as animation 11 | 12 | 13 | np.random.seed(1) 14 | tf.random.set_seed(1) 15 | 16 | 17 | f = open("AC.pckl", "rb") 18 | data = pickle.load(f) 19 | f.close() 20 | 21 | t = np.array(data["t"]).flatten()[:, None] 22 | x = np.array(data["x"]).flatten()[:, None] 23 | Exact_u = np.array(data["u"]) 24 | 25 | ## Positive half of the problem 26 | # x = x[256:] 27 | # Exact_u = Exact_u[256:] 28 | 29 | uxn = len(x) 30 | utn = len(t) 31 | 32 | 33 | lb = np.array([-1.0]) 34 | ub = np.array([1.0]) 35 | 36 | N0 = 256 37 | N_b = 200 38 | N_f = 500 39 | 40 | data_size = 10 41 | 42 | 43 | idx_x = np.random.choice(x.shape[0], N0, replace=False) 44 | x0 = x[idx_x, :] 45 | u0 = Exact_u[idx_x, 0:1] 46 | 47 | # u0 = tf.cast(u0, tf.float3) 48 | idx_t = np.random.choice(t.shape[0], N_b, replace=False) 49 | tb = t[idx_t, :] 50 | 51 | x_f_idx = np.random.choice(x.shape[0], N_f, replace=True) 52 | t_f_idx = np.random.choice(t.shape[0], N_f, replace=True) 53 | 54 | x_f = tf.convert_to_tensor(x[x_f_idx, :], dtype=tf.float32) 55 | t_f = tf.convert_to_tensor(t[t_f_idx, :], dtype=tf.float32) 56 | u_f = tf.convert_to_tensor(Exact_u[x_f_idx, t_f_idx], dtype=tf.float32) 57 | 58 | X0 = np.concatenate((x0, 0 * x0), 1) # (x0, 0) 59 | X_lb = np.concatenate((0 * tb + lb[0], tb), 1) # (lb[0], tb) 60 | X_ub = np.concatenate((0 * tb + ub[0], tb), 1) # (ub[0], tb) 61 | 62 | x0 = X0[:, 0:1] 63 | t0 = X0[:, 1:2] 64 | X_0 = tf.cast(tf.concat([x0, t0], 1), dtype=tf.float32) 65 | 66 | x_lb = tf.convert_to_tensor(X_lb[:, 0:1], dtype=tf.float32) 67 | t_lb = tf.convert_to_tensor(X_lb[:, 1:2], dtype=tf.float32) 68 | 69 | x_ub = tf.convert_to_tensor(X_ub[:, 0:1], dtype=tf.float32) 70 | t_ub = tf.convert_to_tensor(X_ub[:, 1:2], dtype=tf.float32) 71 | 72 | 73 | splits = np.linspace( 74 | int(0.15 * uxn), int(0.85 * uxn), data_size, endpoint=False 75 | ).astype(int) 76 | 77 | x_d = x[splits, :][:, None] 78 | X_D, T_D = np.meshgrid(x_d, t) 79 | X_D = np.hstack((X_D.flatten()[:, None], T_D.flatten()[:, None])) 80 | 81 | u_d = Exact_u[splits] 82 | u_d_flat = u_d.T.flatten()[:, None] 83 | 84 | 85 | def f_model(w, b, x, t): 86 | # keep track of our gradients 87 | with tf.GradientTape(persistent=True) as tape: 88 | tape.watch(x) 89 | tape.watch(t) 90 | 91 | # Getting the prediction 92 | u = multilayer_perceptron(w, b, tf.concat([x, t], 1)) 93 | # _, u_x = u_x_model(x, t) 94 | u_x = tape.gradient(u, x) 95 | # Getting the other derivatives 96 | u_xx = tape.gradient(u_x, x) 97 | u_t = tape.gradient(u, t) 98 | 99 | # Letting the tape go 100 | del tape 101 | 102 | f_u = u_t - 0.0001 * u_xx + 5.0 * u**3 - 5.0 * u 103 | 104 | return f_u 105 | 106 | 107 | def u_x_model(w, b, x, t): 108 | with tf.GradientTape(persistent=True) as tape: 109 | tape.watch(x) 110 | tape.watch(t) 111 | X = tf.concat([x, t], 1) 112 | u = multilayer_perceptron(w, b, X) 113 | # u = u[:, 0:1] 114 | 115 | u_x = tape.gradient(u, x) 116 | 117 | del tape 118 | 119 | return u, u_x 120 | 121 | 122 | @tf.function 123 | def loss(w, b): 124 | 125 | f_u_pred = f_model(w, b, x_f, t_f) 126 | 127 | u0_pred = multilayer_perceptron(w, b, X_0) 128 | u_lb_pred, u_x_lb_pred = u_x_model(w, b, x_lb, t_lb) 129 | u_ub_pred, u_x_ub_pred = u_x_model(w, b, x_ub, t_ub) 130 | 131 | u_d_pred = multilayer_perceptron(w, b, X_D.astype(np.float32)) 132 | 133 | mse_u = tf.reduce_mean(tf.square(u0 - u0_pred)) 134 | 135 | mse_b = tf.reduce_mean(tf.square(u_lb_pred - u_ub_pred)) + tf.reduce_mean( 136 | tf.square(u_x_lb_pred + u_x_ub_pred) 137 | ) 138 | 139 | mse_f = tf.reduce_mean(tf.square(f_u_pred)) 140 | 141 | mse_d = tf.reduce_mean(tf.square(u_d_flat - u_d_pred)) 142 | 143 | # return 100 * mse_u + mse_b + mse_f + 20 * mse_d 144 | return 10 * mse_u + 10 * mse_d + mse_f + mse_b 145 | 146 | 147 | def loss_grad(): 148 | def _loss(w, b): 149 | with tf.GradientTape(persistent=True) as tape: 150 | tape.watch(w) 151 | tape.watch(b) 152 | loss_value = loss(w, b) 153 | trainable_variables = w + b 154 | grads = tape.gradient(loss_value, trainable_variables) 155 | return loss_value, grads 156 | 157 | return _loss 158 | 159 | 160 | def run_swarm(swarm, X): 161 | swarm_y = [] 162 | for particle in swarm: 163 | w, b = decode(particle, layer_sizes) 164 | swarm_y.append(multilayer_perceptron(w, b, X)) 165 | return swarm_y 166 | 167 | 168 | def format_time(seconds): 169 | m, s = divmod(seconds, 60) 170 | h, m = divmod(m, 60) 171 | return h, m, s 172 | 173 | 174 | layer_sizes = [2] + 5 * [15] + [1] 175 | pop_size = 10 176 | n_iter = 6000 177 | 178 | opt = pso( 179 | loss_grad(), 180 | layer_sizes, 181 | n_iter, 182 | pop_size, 183 | 0.999, 184 | 8e-4, 185 | 5e-3, 186 | initialization_method="xavier", 187 | verbose=True, 188 | gd_alpha=1e-4, 189 | ) 190 | 191 | start = time.time() 192 | opt.train() 193 | end = time.time() 194 | print("Time elapsed : %2d:%2d:%2d" % format_time(end - start)) 195 | 196 | 197 | X, T = np.meshgrid(x, t) 198 | X_flat = np.hstack((X.flatten()[:, None], T.flatten()[:, None])) 199 | u_star = Exact_u.T.flatten() 200 | x_ = np.squeeze(x) 201 | swarm = opt.get_swarm() 202 | preds = run_swarm(swarm, X_flat.astype(np.float32)) 203 | mean = tf.squeeze(tf.reduce_mean(preds, axis=0)) 204 | var = tf.squeeze(tf.math.reduce_std(preds, axis=0)) 205 | print("Last Loss: ", opt.loss_history[-1]) 206 | 207 | time_steps = utn - 1 # total number of time steps in animation 208 | fps = 15 # frames/second of animation 209 | error_u = np.linalg.norm(u_star - mean, 2) / np.linalg.norm(u_star, 2) 210 | print("Error u: %e" % (error_u)) 211 | 212 | 213 | def snapshot(i): 214 | l_ind = i * uxn 215 | u_ind = (i + 1) * uxn 216 | plt.clf() 217 | for k in range(len(preds)): 218 | plt.plot(x_, preds[k][l_ind:u_ind], linewidth=0.3) 219 | plt.scatter(x_d, u_d_flat[i * data_size : (i + 1) * data_size]) 220 | plt.plot(x_, u_star[l_ind:u_ind], "b-", linewidth=3, label="Exact") 221 | plt.plot( 222 | x_, 223 | mean[l_ind:u_ind], 224 | "r--", 225 | linewidth=3, 226 | label=opt.name, 227 | ) 228 | plt.fill_between( 229 | x_, 230 | mean[l_ind:u_ind] - 3 * var[l_ind:u_ind], 231 | mean[l_ind:u_ind] + 3 * var[l_ind:u_ind], 232 | color="gray", 233 | alpha=0.2, 234 | ) 235 | 236 | plt.xlabel("$x$", fontsize="xx-large") 237 | plt.ylabel("$u(t,x)$", fontsize="xx-large") 238 | plt.xlim(-1.0, 1.0) 239 | plt.ylim(-1.02, 1.02) 240 | plt.grid() 241 | plt.legend(fontsize="x-large") 242 | 243 | 244 | fig = plt.figure(figsize=(8, 8), dpi=150) 245 | # Call the animator: 246 | anim = animation.FuncAnimation(fig, snapshot, frames=time_steps) 247 | # Save the animation as an mp4. This requires ffmpeg to be installed. 248 | anim.save("ac_demo.gif", fps=fps) 249 | -------------------------------------------------------------------------------- /src/swarm/optimizers/pso.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from numpy.random import uniform 3 | from swarm import utils 4 | from numpy import linspace 5 | 6 | 7 | class pso: 8 | def __init__( 9 | self, 10 | loss_op, 11 | layer_sizes, 12 | n_iter=2000, 13 | pop_size=30, 14 | b=0.9, 15 | c1=0.8, 16 | c2=0.5, 17 | x_min=-1, 18 | x_max=1, 19 | cold_start=True, 20 | initialization_method=None, 21 | verbose=False, 22 | c_decrease=False, 23 | pre_trained_x=None, 24 | ): 25 | """The Particle Swarm Optimizer class. Specially built to deal with tensorflow neural networks. 26 | 27 | Args: 28 | loss_op (function): The fitness function for PSO. 29 | layer_sizes (list): The layers sizes of the neural net. 30 | n_iter (int, optional): Number of PSO iterations. Defaults to 2000. 31 | pop_size (int, optional): Population of the PSO swarm. Defaults to 30. 32 | b (float, optional): Inertia of the particles. Defaults to 0.9. 33 | c1 (float, optional): The *p-best* coeficient. Defaults to 0.8. 34 | c2 (float, optional): The *g-best* coeficient. Defaults to 0.5. 35 | x_min (int, optional): The min value for the weights generation. Defaults to -1. 36 | x_max (int, optional): The max value for the weights generation. Defaults to 1. 37 | cold_start (bool, optional): Set the starting velocities to 0. Defaults to True. 38 | initialization_method (_type_, optional): Chooses how to initialize the Neural Net weights. Allowed to be one of "uniform", "xavier", or "log_logistic". Defaults to None, where it uses uniform initialization. 39 | verbose (bool, optional): Shows info during the training . Defaults to False. 40 | """ 41 | self.loss_op = loss_op 42 | self.layer_sizes = layer_sizes 43 | self.pop_size = pop_size 44 | self.dim = utils.dimensions(layer_sizes) 45 | self.n_iter = n_iter 46 | self.b = b 47 | self.c1 = c1 48 | self.c2 = c2 49 | self.x_min = x_min 50 | self.x_max = x_max 51 | self.initialization_method = initialization_method 52 | self.x = ( 53 | pre_trained_x if pre_trained_x is not None else self.build_swarm() 54 | ) 55 | self.p = self.x 56 | self.loss_history = [] 57 | self.f_p, self.grads = self.fitness_fn(self.p) 58 | self.g = self.p[tf.math.argmin(input=self.f_p).numpy()[0]] 59 | self.cold_start = cold_start 60 | self.v = self.start_velocities() 61 | self.verbose = verbose 62 | self.name = "PSO" 63 | self.c_decrease = c_decrease 64 | self.verbose_milestone = linspace(0, n_iter, 11).astype(int) 65 | 66 | def build_swarm(self): 67 | """Creates the swarm following the selected initialization method. 68 | 69 | Args: 70 | initialization_method (str): Chooses how to initialize the Neural Net weights. Allowed to be one of "uniform", "xavier", or "log_logistic". Defaults to None, where it uses uniform initialization. 71 | 72 | Returns: 73 | tf.Tensor: The PSO swarm population. Each particle represents a neural network. 74 | """ 75 | return utils.build_NN( 76 | self.pop_size, self.layer_sizes, self.initialization_method 77 | ) 78 | 79 | def update_pso_params(self): 80 | self.c1 = self.c1 - self.c1 / self.n_iter 81 | self.c2 = self.c2 - self.c2 / self.n_iter 82 | 83 | def start_velocities(self): 84 | """Start the velocities of each particle in the population (swarm). If 'self.cold_start' is 'TRUE', the swarm starts with velocity 0, which means stopped. 85 | 86 | Returns: 87 | tf.Tensor: The starting velocities. 88 | """ 89 | if self.cold_start: 90 | return tf.zeros([self.pop_size, self.dim]) 91 | else: 92 | return tf.Variable( 93 | tf.random.uniform( 94 | [self.pop_size, self.dim], 95 | -self.x_max - self.x_min, 96 | self.x_max - self.x_min, 97 | ) 98 | ) 99 | 100 | def individual_fn(self, particle): 101 | """Auxiliary function to get the loss of each particle. 102 | 103 | Args: 104 | particle (tf.Tensor): One particle of the PSO swarm representing a full neural network. 105 | 106 | Returns: 107 | tuple: The loss value and the gradients. 108 | """ 109 | w, b = utils.decode(particle, self.layer_sizes) 110 | loss, grad = self.loss_op(w, b) 111 | return loss, utils.flat_grad(grad) 112 | 113 | @tf.function 114 | def fitness_fn(self, x): 115 | """Fitness function for the whole swarm. 116 | 117 | Args: 118 | x (tf.Tensor): The swarm. All the particle's current positions. Which means the weights of all neural networks. 119 | 120 | Returns: 121 | tuple: the losses and gradients for all particles. 122 | """ 123 | f_x, grads = tf.vectorized_map(self.individual_fn, x) 124 | return f_x[:, None], grads 125 | 126 | def get_randoms(self): 127 | """Generate random values to update the particles' positions. 128 | 129 | Returns: 130 | _type_: tf.Tensor 131 | """ 132 | return uniform(0, 1, [2, self.dim])[:, None] 133 | 134 | def update_p_best(self): 135 | """Updates the *p-best* positions.""" 136 | f_x, self.grads = self.fitness_fn(self.x) 137 | self.loss_history.append(tf.reduce_mean(f_x).numpy()) 138 | self.p = tf.where(f_x < self.f_p, self.x, self.p) 139 | self.f_p = tf.where(f_x < self.f_p, f_x, self.f_p) 140 | 141 | def update_g_best(self): 142 | """Update the *g-best* position.""" 143 | self.g = self.p[tf.math.argmin(input=self.f_p).numpy()[0]] 144 | 145 | def step(self): 146 | """It runs ONE step on the particle swarm optimization.""" 147 | r1, r2 = self.get_randoms() 148 | self.v = ( 149 | self.b * self.v 150 | + self.c1 * r1 * (self.p - self.x) 151 | + self.c2 * r2 * (self.g - self.x) 152 | ) 153 | self.x = self.x + self.v 154 | self.update_p_best() 155 | self.update_g_best() 156 | 157 | def train(self): 158 | """The particle swarm optimization. The PSO will optimize the weights according to the losses of the neural network, so this process is actually the neural network training.""" 159 | for i in range(self.n_iter): 160 | self.step() 161 | if self.c_decrease: 162 | self.update_pso_params() 163 | if self.verbose and i in self.verbose_milestone: 164 | utils.progress( 165 | (i / self.n_iter) * 100, 166 | metric="loss", 167 | metricValue=self.loss_history[-1], 168 | ) 169 | if self.verbose: 170 | utils.progress(100) 171 | print() 172 | 173 | def get_best(self): 174 | """Return the *g-best*, the particle with best results after the training. 175 | 176 | Returns: 177 | tf.Tensor: the best particle of the swarm. 178 | """ 179 | return utils.decode(self.g, self.layer_sizes) 180 | 181 | def get_swarm(self): 182 | """Return the swarm. 183 | 184 | Returns: 185 | tf.Tensor: The positions of each particle. 186 | """ 187 | return self.x 188 | 189 | def set_n_iter(self, n_iter): 190 | """Set the number of iterations. 191 | Args: 192 | x (int): Number of iterations. 193 | """ 194 | self.n_iter = n_iter 195 | self.verbose_milestone = linspace(0, n_iter, 11).astype(int) 196 | -------------------------------------------------------------------------------- /src/swarm/optimizers/pso_adam.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from numpy import inf, linspace 3 | from numpy.random import uniform 4 | from swarm import utils 5 | 6 | 7 | class pso: 8 | def __init__( 9 | self, 10 | loss_op, 11 | layer_sizes, 12 | n_iter=2000, 13 | pop_size=30, 14 | b=0.99, 15 | c1=8e-2, 16 | c2=5e-1, 17 | gd_alpha=5e-3, 18 | initialization_method=None, 19 | verbose=False, 20 | c_decrease=False, 21 | pre_trained_x=None, 22 | beta_1=0.99, 23 | beta_2=0.999, 24 | epsilon=1e-8, 25 | ): 26 | """The Particle Swarm Optimizer class. Specially built to deal with tensorflow neural networks. 27 | 28 | Args: 29 | loss_op (function): The fitness function for PSO. 30 | layer_sizes (list): The layers sizes of the neural net. 31 | n_iter (int, optional): Number of PSO iterations. Defaults to 2000. 32 | pop_size (int, optional): Population of the PSO swarm. Defaults to 30. 33 | b (float, optional): Inertia of the particles. Defaults to 0.9. 34 | c1 (float, optional): The *p-best* coeficient. Defaults to 0.8. 35 | c2 (float, optional): The *g-best* coeficient. Defaults to 0.5. 36 | gd_alpha (float, optional): Learning rate for gradient descent. Defaults to 0.00, so there wouldn't have any gradient-based optimization. 37 | initialization_method (_type_, optional): Chooses how to initialize the Neural Net weights. Allowed to be one of "uniform", "xavier", or "log_logistic". Defaults to None, where it uses uniform initialization. 38 | verbose (bool, optional): Shows info during the training . Defaults to False. 39 | """ 40 | self.loss_op = loss_op 41 | self.layer_sizes = layer_sizes 42 | self.pop_size = pop_size 43 | self.dim = utils.dimensions(layer_sizes) 44 | self.n_iter = n_iter 45 | self.b = b 46 | self.c1 = c1 47 | self.c2 = c2 48 | self.initialization_method = initialization_method 49 | self.x = ( 50 | pre_trained_x if pre_trained_x is not None else self.build_swarm() 51 | ) 52 | self.f_x, self.grads = self.fitness_fn(self.x) 53 | self.p, self.f_p = self.x, self.f_x 54 | self.loss_history = [] 55 | self.g = self.p[tf.math.argmin(input=self.f_p).numpy()[0]] 56 | self.v = self.start_velocities() 57 | self.verbose = verbose 58 | self.c_decrease = c_decrease 59 | self.verbose_milestone = linspace(0, n_iter, 11).astype(int) 60 | self.beta_1 = beta_1 61 | self.beta_2 = beta_2 62 | self.epsilon = (epsilon,) 63 | self.m1 = tf.zeros([self.pop_size, self.dim]) 64 | self.m2 = tf.zeros([self.pop_size, self.dim]) 65 | self.gd_alpha = ( 66 | gd_alpha * tf.math.sqrt(1 - self.beta_2) / (1 - self.beta_1) 67 | ) 68 | self.name = "PSO" if self.gd_alpha == 0 else "PSO-GD" 69 | 70 | def build_swarm(self): 71 | """Creates the swarm following the selected initialization method. 72 | 73 | Args: 74 | initialization_method (str): Chooses how to initialize the Neural Net weights. Allowed to be one of "uniform", "xavier", or "log_logistic". Defaults to None, where it uses uniform initialization. 75 | 76 | Returns: 77 | tf.Tensor: The PSO swarm population. Each particle represents a neural network. 78 | """ 79 | return utils.build_NN( 80 | self.pop_size, self.layer_sizes, self.initialization_method 81 | ) 82 | 83 | def update_pso_params(self): 84 | self.c1 = self.c1 - 2 * self.c1 / self.n_iter 85 | self.c2 = self.c2 + self.c2 / self.n_iter 86 | 87 | def start_velocities(self): 88 | """Start the velocities of each particle in the population (swarm) as `0`. 89 | 90 | Returns: 91 | tf.Tensor: The starting velocities. 92 | """ 93 | return tf.zeros([self.pop_size, self.dim]) 94 | 95 | def individual_fn(self, particle): 96 | """Auxiliary function to get the loss of each particle. 97 | 98 | Args: 99 | particle (tf.Tensor): One particle of the PSO swarm representing a full neural network. 100 | 101 | Returns: 102 | tuple: The loss value and the gradients. 103 | """ 104 | w, b = utils.decode(particle, self.layer_sizes) 105 | loss, grad = self.loss_op(w, b) 106 | return loss, utils.flat_grad(grad) 107 | 108 | @tf.function 109 | def fitness_fn(self, x): 110 | """Fitness function for the whole swarm. 111 | 112 | Args: 113 | x (tf.Tensor): The swarm. All the particle's current positions. Which means the weights of all neural networks. 114 | 115 | Returns: 116 | tuple: the losses and gradients for all particles. 117 | """ 118 | f_x, grads = tf.vectorized_map(self.individual_fn, x) 119 | return f_x[:, None], grads 120 | 121 | def get_randoms(self): 122 | """Generate random values to update the particles' positions. 123 | 124 | Returns: 125 | _type_: tf.Tensor 126 | """ 127 | return uniform(0, 1, [2, self.dim])[:, None] 128 | 129 | def update_p_best(self): 130 | """Updates the *p-best* positions.""" 131 | self.p = tf.where(self.f_x < self.f_p, self.x, self.p) 132 | self.f_p = tf.where(self.f_x < self.f_p, self.f_x, self.f_p) 133 | 134 | def update_g_best(self): 135 | """Update the *g-best* position.""" 136 | self.g = self.p[tf.math.argmin(input=self.f_p).numpy()[0]] 137 | 138 | def gradient_descent(self): 139 | self.m1 = self.beta_1 * self.m1 + (1 - self.beta_1) * self.grads 140 | self.m2 = self.beta_2 * self.m2 + (1 - self.beta_2) * tf.math.square( 141 | self.grads 142 | ) 143 | return self.gd_alpha * self.m1 / tf.math.sqrt(self.m2) + self.epsilon 144 | 145 | def step(self): 146 | """It runs ONE step on the particle swarm optimization.""" 147 | r1, r2 = self.get_randoms() 148 | self.v = self.b * self.v + (1 - self.b) * ( 149 | self.c1 * r1 * (self.p - self.x) + self.c2 * r2 * (self.g - self.x) 150 | ) 151 | self.x = self.x + self.v - self.gradient_descent() 152 | self.f_x, self.grads = self.fitness_fn(self.x) 153 | self.update_p_best() 154 | self.update_g_best() 155 | 156 | def train(self): 157 | """The particle swarm optimization. The PSO will optimize the weights according to the losses of the neural network, so this process is actually the neural network training.""" 158 | for i in range(self.n_iter): 159 | self.step() 160 | self.loss_history.append(tf.reduce_mean(self.f_x).numpy()) 161 | if self.c_decrease: 162 | self.update_pso_params() 163 | if self.verbose and i in self.verbose_milestone: 164 | utils.progress( 165 | (i / self.n_iter) * 100, 166 | metric="loss", 167 | metricValue=self.loss_history[-1], 168 | ) 169 | if self.verbose: 170 | utils.progress(100) 171 | print() 172 | 173 | def get_best(self): 174 | """Return the *g-best*, the particle with best results after the training. 175 | 176 | Returns: 177 | tf.Tensor: the best particle of the swarm. 178 | """ 179 | return utils.decode(self.g, self.layer_sizes) 180 | 181 | def get_swarm(self): 182 | """Return the swarm. 183 | 184 | Returns: 185 | tf.Tensor: The positions of each particle. 186 | """ 187 | return self.x 188 | 189 | def set_n_iter(self, n_iter): 190 | """Set the number of iterations. 191 | Args: 192 | x (int): Number of iterations. 193 | """ 194 | self.n_iter = n_iter 195 | self.verbose_milestone = linspace(0, n_iter, 11).astype(int) 196 | -------------------------------------------------------------------------------- /src/swarm/utils.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import tensorflow_probability as tfp 3 | from swarm import utils 4 | 5 | 6 | def dimensions(layers): 7 | """Returns the number of `weights + biases` of a neural net, given the layers. 8 | 9 | Args: 10 | layers (list): List of 'int' representing the size of each layer. 11 | 12 | Returns: 13 | int: The number of optimizable variables (`weights + biases`) 14 | """ 15 | sum = 0 16 | for l in range(0, len(layers) - 1): 17 | sum += layers[l] * layers[l + 1] + layers[l + 1] 18 | return sum 19 | 20 | 21 | def progress(percent=0, width=30, metric=None, metricValue=None): 22 | """Prints on screen the current progress of a process. 23 | 24 | Args: 25 | percent (int, optional): The current progress (in percentiles). Defaults to 0. 26 | width (int, optional): The width size of the progress bar. Defaults to 30. 27 | metric (float, optional): The metric used. Defaults to None. 28 | metricValue (str, optional): The unit name for the metric used. Defaults to None. 29 | """ 30 | left = width * int(percent) // 100 31 | right = width - left 32 | if metric: 33 | print( 34 | "\r[", 35 | "#" * left, 36 | " " * right, 37 | "]", 38 | f" {percent:.0f}% -- ", 39 | "Current ", 40 | metric, 41 | ": ", 42 | metricValue, 43 | sep="", 44 | end="", 45 | flush=True, 46 | ) 47 | else: 48 | print( 49 | "\r[", 50 | "#" * left, 51 | " " * right, 52 | "]", 53 | f" {percent:.0f}% -- ", 54 | sep="", 55 | end="", 56 | flush=True, 57 | ) 58 | 59 | 60 | def encode(weights, biases): 61 | """Given weights and biases of a neural net, it returns a flat `tf.Tensor` with those values. This *encoded* format represents a particle in the PSO. 62 | 63 | Args: 64 | weights (tf.Tensor): The weights of the neural net. 65 | biases (tf.Tensor): The biases of the neural net. 66 | 67 | Returns: 68 | tf.Tensor: The particle for PSO. The flattened tensor with weights and biases. 69 | """ 70 | encoded = tf.Variable([]) 71 | for l in weights: 72 | encoded = tf.concat(values=[encoded, tf.reshape(l, -1)], axis=-1) 73 | for l in biases: 74 | encoded = tf.concat(values=[encoded, tf.reshape(l, -1)], axis=-1) 75 | return encoded 76 | 77 | 78 | def decode(encoded, layers): 79 | """It will decode a PSO particle into the weights and biases of a neural network. It does the inverse process of the `encode` function. 80 | 81 | Args: 82 | encoded (tf.Tensor): The PSO particle. 83 | layers (list): List of 'int' representing the size of each layer. 84 | 85 | Returns: 86 | tuple: Two `tf.Tensor` representing the weights and biases of a neural net. 87 | """ 88 | weights = [] 89 | biases = [] 90 | last_cut = 0 91 | num_layers = len(layers) 92 | for l in range(0, num_layers - 1): 93 | next_cut = layers[l] * layers[l + 1] 94 | W = tf.reshape( 95 | tf.slice(encoded, [last_cut], [next_cut]), [layers[l], layers[l + 1]] 96 | ) 97 | last_cut = last_cut + next_cut 98 | weights.append(W) 99 | for l in range(1, num_layers): 100 | b = tf.slice(encoded, [last_cut], [layers[l]]) 101 | last_cut += layers[l] 102 | biases.append(b) 103 | return weights, biases 104 | 105 | 106 | def multilayer_perceptron(weights, biases, X, x_min=-1, x_max=1): 107 | """It runs the multilayer perceptron neural network. Given the weights and biases representing the neural net and the input population `X`. 108 | 109 | Args: 110 | weights (tf.Tensor): The weights of the neural net. 111 | biases (tf.Tensor): The biases of the neural net. 112 | X (tf.Tensor): The input values. 113 | x_min (int, optional): The floor value for the normalization. Defaults to -1. 114 | x_max (int, optional): The roof value for the normalization. Defaults to 1. 115 | 116 | Returns: 117 | tf.Tensor: The prediction `Y`. 118 | """ 119 | num_layers = len(weights) + 1 120 | H = 2.0 * (X - x_min) / (x_max - x_min) - 1.0 121 | for l in range(0, num_layers - 2): 122 | W = weights[l] 123 | b = biases[l] 124 | H = tf.nn.tanh(tf.add(tf.matmul(H, W), b)) 125 | W = weights[-1] 126 | b = biases[-1] 127 | Y = tf.add(tf.matmul(H, W), b) 128 | return Y 129 | 130 | 131 | def replacenan(t): 132 | """Replace `nan` with zeros. **CAUTION**: `nan` may be the result of an infinitely small number, but it could happen the other way around too. If the `nan` was the result of an infinitely big number, the zero representation would be misleading. 133 | 134 | Args: 135 | t (tf.Tensor): The tensor with `nan` values. 136 | 137 | Returns: 138 | tf.Tensor: Tensor with `0s` instead of `nan`. 139 | """ 140 | return tf.where(tf.math.is_nan(t), tf.zeros_like(t), t) 141 | 142 | 143 | def layer_init(size, method): 144 | """Initialization for normalized a layer. 145 | 146 | Args: 147 | size (int): The layer size. 148 | 149 | Returns: 150 | tf.Tensor: The weights for the layer. 151 | """ 152 | in_dim = size[0] 153 | out_dim = size[1] 154 | _stddev = tf.sqrt(6 / (in_dim + out_dim)) 155 | if method == "he": 156 | _stddev = tf.sqrt(2 / (in_dim)) 157 | if method == "lecun": 158 | _stddev = tf.sqrt(1 / (in_dim)) 159 | return tf.Variable( 160 | tf.random.truncated_normal([in_dim, out_dim], stddev=_stddev), 161 | dtype=tf.float32, 162 | ) 163 | 164 | 165 | def initialize_NN(layers, method): 166 | """Initialize a neural network following the initialization given by `method`. 167 | 168 | Args: 169 | layers (list): A list of `int` representing each layer size. 170 | 171 | Returns: 172 | tuple: Two `tf.Tensor` with the weights and biases of the neural net. 173 | """ 174 | weights = [] 175 | biases = [] 176 | num_layers = len(layers) 177 | for l in range(0, num_layers - 1): 178 | W = layer_init([layers[l], layers[l + 1]], method) 179 | b = tf.Variable( 180 | tf.zeros([1, layers[l + 1]], dtype=tf.float32), dtype=tf.float32 181 | ) 182 | weights.append(W) 183 | biases.append(b) 184 | return weights, biases 185 | 186 | 187 | def _build_normalized(pop_size, layer_sizes, method): 188 | """Initialize multiple neural networks using a normalized initialization selected by `method`. 189 | 190 | Args: 191 | pop_size (int): Number of neural networks to initialize. 192 | layers (list): A list of `int` representing each layer size. (All the neural nets must have the same topology). 193 | 194 | Returns: 195 | tf.Tensor: All the neural nets. 196 | """ 197 | init_nns = [] 198 | for _ in range(pop_size): 199 | w, b = initialize_NN(layer_sizes, method) 200 | new_nn = encode(w, b) 201 | init_nns.append(new_nn) 202 | return tf.Variable(init_nns, dtype=tf.float32) 203 | 204 | 205 | def _build_uniform(pop_size, layer_sizes, _): 206 | """Initialize multiple neural networks using a uniformly distributed initialization. 207 | 208 | Args: 209 | pop_size (int): Number of neural networks to initialize. 210 | layers (list): A list of `int` representing each layer size. (All the neural nets must have the same topology). 211 | 212 | Returns: 213 | tf.Tensor: All the neural nets. 214 | """ 215 | x_min = -1 216 | x_max = 1 217 | dim = dimensions(layer_sizes) 218 | return tf.Variable(tf.random.uniform([pop_size, dim], x_min, x_max)) 219 | 220 | 221 | def _build_logLogistic(pop_size, layer_sizes, _): 222 | """Initialize multiple neural networks using a log-logistic initialization. 223 | 224 | Args: 225 | pop_size (int): Number of neural networks to initialize. 226 | layers (list): A list of `int` representing each layer size. (All the neural nets must have the same topology). 227 | 228 | Returns: 229 | tf.Tensor: All the neural nets. 230 | """ 231 | dim = dimensions(layer_sizes) 232 | dist = tfp.distributions.LogLogistic(0, 0.1) 233 | return dist.sample([pop_size, dim]) 234 | 235 | 236 | def build_NN(pop_size, layer_sizes, method): 237 | """Initialize multiple neural networks using a normalized initialization selected by `method`. 238 | 239 | Args: 240 | pop_size (int): Number of neural networks to initialize. 241 | layers (list): A list of `int` representing each layer size. (All the neural nets must have the same topology). 242 | 243 | Returns: 244 | tf.Tensor: All the neural nets. 245 | """ 246 | methods = { 247 | "uniform": "_build_uniform", 248 | "log_logistic": "_build_logLogistic", 249 | "xavier": "_build_normalized", 250 | "he": "_build_normalized", 251 | "lecun": "_build_normalized", 252 | } 253 | return getattr(utils, methods.get(method, "_build_normalized"))( 254 | pop_size, layer_sizes, method 255 | ) 256 | 257 | 258 | def flat_grad(grad): 259 | """Flattens the gradient tensor. 260 | 261 | Args: 262 | grad (tf.Tensor): Gradients 263 | 264 | Returns: 265 | tf.Tensor: Flatted gradients. 266 | """ 267 | flatted = [] 268 | for g in grad: 269 | flatted.append(tf.reshape(g, [-1])) 270 | return tf.concat(flatted, 0) 271 | 272 | 273 | def dominance(x, y, weak=False): 274 | """Dominance test. True means x dominates y. 275 | 276 | Args: 277 | x (tf.Tensor): A Tensor. Must be one of the following types: float32, float64. 278 | y (tf.Tensor): A Tensor. Must have the same type as x. 279 | weak (bool, optional): True for the weak dominance. Defaults to False. 280 | 281 | Returns: 282 | tf.Tensor: A Tensor of type bool. 283 | """ 284 | less_equal = tf.reduce_all(tf.math.less_equal(x, y), 1) 285 | if weak: 286 | return less_equal 287 | less = tf.reduce_any(tf.math.less(x, y), 1) 288 | return tf.logical_and(less_equal, less) 289 | 290 | 291 | def normalize(arr, t=[-1, 1]): 292 | """Min-Max normalization 293 | 294 | Args: 295 | arr (tf.Tensor): A Tensor. Must be one of the following types: float32, float64. 296 | t (list, optional): Range for the final tensor. Defaults to [-1, 1]. 297 | 298 | Returns: 299 | tf.Tensor: A Tensor of same type as x. 300 | """ 301 | max = tf.reduce_max(arr) 302 | min = tf.reduce_min(arr) 303 | return t[0] + (arr - min) * (t[1] - t[0]) / (max - min + 1e-8) 304 | --------------------------------------------------------------------------------