├── .gitignore ├── LICENSE ├── README.md ├── data ├── err_dirichlet.png ├── err_neumann.png ├── flash_noflash │ ├── Objects_010_ambient.png │ ├── Objects_010_flash.png │ ├── Rooms_010_ambient.png │ ├── Rooms_010_flash.png │ ├── Shelves_007_ambient.png │ ├── Shelves_007_flash.png │ ├── Shelves_009_ambient.png │ ├── Shelves_009_flash.png │ ├── Toys_007_ambient.png │ ├── Toys_007_flash.png │ ├── artifacts.png │ └── result.png ├── gt_dirichlet.png ├── gt_neumann.png ├── laplace_sol_2d.png ├── laplace_sol_3d.png ├── poisson_image_editing │ ├── mask.jpg │ ├── result.png │ ├── source.jpg │ └── target.jpg ├── sol_dirichlet.png ├── sol_neumann.png ├── zero_mean_false.png └── zero_mean_true.png ├── examples ├── context.py ├── flash_noflash.py ├── numpy_mode.py ├── poisson_image_editing.py ├── sympy_dirichlet.py ├── sympy_laplace_eq.py ├── sympy_neumann.py ├── sympy_region_laplace.py └── sympy_zero_mean.py ├── poissonpy ├── __init__.py ├── functional.py ├── helpers.py ├── solvers.py └── utils.py ├── requirements.txt └── setup.py /.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/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | #.idea/ -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # poissonpy 2 | Plug-and-play standalone library for solving 2D Poisson equations. Useful tool in scientific computing prototyping, image and video processing, computer graphics. 3 | 4 | ## Features 5 | - Solves the Poisson equation on sqaure or non-square rectangular grids. 6 | - Solves the Poisson equation on regions with arbitrary shape. 7 | - Supports arbitrary boundary and interior conditions using `sympy` function experssions or `numpy` arrays. 8 | - Supports Dirichlet, Neumann, or mixed boundary conditions. 9 | 10 | ## Disclaimer 11 | This package is only used to solve 2D Poisson equations. If you are looking for a general purpose and optimized PDE library, you might want to checkout the [FEniCSx project](https://fenicsproject.org/index.html). 12 | 13 | ## Usage 14 | Import necessary libraries. `poissonpy` utilizes `numpy` and `sympy` greatly, so its best to import both: 15 | 16 | ```python 17 | import numpy as np 18 | from sympy import sin, cos 19 | from sympy.abc import x, y 20 | 21 | from poissonpy import functional, utils, sovlers 22 | ``` 23 | 24 | ### Defining `sympy` functions 25 | In the following examples, we use a ground truth function to create a mock Poisson equation and compare the solver's solution with the analytical solution. 26 | 27 | Define functions using `sympy` function expressions or `numpy` arrays: 28 | 29 | ```python 30 | f_expr = sin(x) + cos(y) # create sympy function expression 31 | laplacian_expr = functional.get_sp_laplacian_expr(f_expr) # create sympy laplacian function expression 32 | 33 | f = functional.get_sp_function(f_expr) # create sympy function 34 | laplacian = functional.get_sp_function(laplacian_expr) # create sympy function 35 | ``` 36 | 37 | ### Dirichlet Boundary Conditions 38 | Define interior and Dirichlet boundary conditions: 39 | 40 | ```python 41 | interior = laplacian 42 | boundary = { 43 | "left": (f, "dirichlet"), 44 | "right": (f, "dirichlet"), 45 | "top": (f, "dirichlet"), 46 | "bottom": (f, "dirichlet") 47 | } 48 | ``` 49 | 50 | Initialize solver and solve Poisson equation: 51 | 52 | ```python 53 | solver = Poisson2DRectangle(((-2*np.pi, -2*np.pi), (2*np.pi, 2*np.pi)), 54 | interior, boundary, X=100, Y=100) 55 | solution = solver.solve() 56 | ``` 57 | 58 | Plot solution and ground truth: 59 | ```python 60 | poissonpy.plot_3d(solver.x_grid, solver.y_grid, solution) 61 | poissonpy.plot_3d(solver.x_grid, solver.y_grid, f(solver.x_grid, solver.y_grid)) 62 | ``` 63 | 64 | |Solution|Ground truth|Error| 65 | |--|--|--| 66 | |![](data/sol_dirichlet.png)|![](data/gt_dirichlet.png)|![](data/err_dirichlet.png)| 67 | 68 | ### Neumann Boundary Conditions 69 | You can also define Neumann boundary conditions by specifying `neumann_x` and `neumann_y` in the boundary condition parameter. 70 | 71 | ```python 72 | 73 | x_derivative_expr = functional.get_sp_derivative_expr(f_expr, x) 74 | y_derivative_expr = functional.get_sp_derivative_expr(f_expr, y) 75 | 76 | interior = laplacian 77 | boundary = { 78 | "left": (f, "dirichlet"), 79 | "right": (functional.get_sp_function(x_derivative_expr), "neumann_x"), 80 | "top": (f, "dirichlet"), 81 | "bottom": (functional.get_sp_function(y_derivative_expr), "neumann_y") 82 | } 83 | ``` 84 | 85 | |Solution|Ground truth|Error| 86 | |--|--|--| 87 | |![](data/sol_neumann.png)|![](data/gt_neumann.png)|![](data/err_neumann.png)| 88 | 89 | ### Zero-mean solution 90 | If the boundary condition is purely Neumann, then the solution is not unique. Naively solving the Poisson equation gives bad results. In this case, you can set the `zero_mean` paramter to `True`, such that the solver finds a zero-mean solution. 91 | 92 | ```python 93 | solver = solvers.Poisson2DRectangle( 94 | ((-2*np.pi, -2*np.pi), (2*np.pi, 2*np.pi)), interior, boundary, 95 | X=100, Y=100, zero_mean=True) 96 | ``` 97 | 98 | |`zero_mean=False`|`zero_mean=True`|Ground truth| 99 | |--|--|--| 100 | |![](data/zero_mean_false.png)|![](data/zero_mean_true.png)|![](data/gt_neumann.png)| 101 | 102 | ### Laplace Equation 103 | It's also straightforward to define a Laplace equation - **we simply set the interior laplacian value to 0**. In the following example, we set the boundary values to be spatially-varying periodic functions. 104 | 105 | ```python 106 | interior = 0 # laplace equation form 107 | left = poissonpy.get_2d_sympy_function(sin(y)) 108 | right = poissonpy.get_2d_sympy_function(sin(y)) 109 | top = poissonpy.get_2d_sympy_function(sin(x)) 110 | bottom = poissonpy.get_2d_sympy_function(sin(x)) 111 | 112 | boundary = { 113 | "left": (left, "dirichlet"), 114 | "right": (right, "dirichlet"), 115 | "top": (top, "dirichlet"), 116 | "bottom": (bottom, "dirichlet") 117 | } 118 | ``` 119 | 120 | Solve the Laplace equation: 121 | 122 | ```python 123 | solver = Poisson2DRectangle( 124 | ((-2*np.pi, -2*np.pi), (2*np.pi, 2*np.pi)), interior, boundary, 100, 100) 125 | solution = solver.solve() 126 | poissonpy.plot_3d(solver.x_grid, solver.y_grid, solution, "solution") 127 | poissonpy.plot_2d(solution, "solution") 128 | ``` 129 | 130 | |3D surface plot|2D heatmap| 131 | |--|--| 132 | |![](data/laplace_sol_3d.png)|![](data/laplace_sol_2d.png)| 133 | 134 | ### Arbitrary-shaped domain 135 | Use the `Poisson2DRegion` class to solve the Poisson eqaution on a arbitrary-shaped function domain. `poissonpy` can be seamlessly integrated in gradient-domain image processing algorithms. 136 | 137 | The following is an example where `poissonpy` is used to implement the image cloning algorithm proposed in [Poisson Image Editing](https://www.cs.jhu.edu/~misha/Fall07/Papers/Perez03.pdf) by Perez et al., 2003. See `examples/poisson_image_editing.py` for more details. 138 | 139 | ```python 140 | # compute laplacian of interpolation function 141 | Gx_src, Gy_src = functional.get_np_gradient(source) 142 | Gx_target, Gy_target = functional.get_np_gradient(target) 143 | G_src_mag = (Gx_src**2 + Gy_src**2)**0.5 144 | G_target_mag = (Gx_target**2 + Gy_target**2)**0.5 145 | Gx = np.where(G_src_mag > G_target_mag, Gx_src, Gx_target) 146 | Gy = np.where(G_src_mag > G_target_mag, Gy_src, Gy_target) 147 | Gxx, _ = functional.get_np_gradient(Gx, forward=False) 148 | _, Gyy = functional.get_np_gradient(Gy, forward=False) 149 | laplacian = Gxx + Gyy 150 | 151 | # solve interpolation function 152 | solver = solvers.Poisson2DRegion(mask, laplacian, target) 153 | solution = solver.solve() 154 | 155 | # alpha-blend interpolation and target function 156 | blended = mask * solution + (1 - mask) * target 157 | ``` 158 | 159 |

160 | 161 |

162 | 163 | Another example of using `poissonpy` to implement flash artifacts and reflection removal, using the algorithm proposed in [Removing Photography Artifacts using Gradient Projection and Flash-Exposure Sampling](http://www.cs.columbia.edu/cg/pdfs/114-flashReflectionsRaskarSig05.pdf) by Agrawal et al. 2005. See `examples/flash_noflash.py` for more details. 164 | 165 | ```python 166 | Gx_a, Gy_a = functional.get_np_gradient(ambient) 167 | Gx_f, Gy_f = functional.get_np_gradient(flash) 168 | 169 | # gradient projection 170 | t = (Gx_a * Gx_f + Gy_a * Gy_f) / (Gx_a**2 + Gy_a**2 + 1e-8) 171 | Gx_f_proj = t * Gx_a 172 | Gy_f_proj = t * Gy_a 173 | 174 | # compute laplacian (div of gradient) 175 | lap = functional.get_np_div(Gx_f_proj, Gy_f_proj) 176 | 177 | # integrate laplacian field 178 | solver = solvers.Poisson2DRegion(mask, lap, flash) 179 | res = solver.solve() 180 | ``` 181 | 182 |

183 | 184 |

185 | -------------------------------------------------------------------------------- /data/err_dirichlet.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/err_dirichlet.png -------------------------------------------------------------------------------- /data/err_neumann.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/err_neumann.png -------------------------------------------------------------------------------- /data/flash_noflash/Objects_010_ambient.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Objects_010_ambient.png -------------------------------------------------------------------------------- /data/flash_noflash/Objects_010_flash.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Objects_010_flash.png -------------------------------------------------------------------------------- /data/flash_noflash/Rooms_010_ambient.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Rooms_010_ambient.png -------------------------------------------------------------------------------- /data/flash_noflash/Rooms_010_flash.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Rooms_010_flash.png -------------------------------------------------------------------------------- /data/flash_noflash/Shelves_007_ambient.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Shelves_007_ambient.png -------------------------------------------------------------------------------- /data/flash_noflash/Shelves_007_flash.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Shelves_007_flash.png -------------------------------------------------------------------------------- /data/flash_noflash/Shelves_009_ambient.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Shelves_009_ambient.png -------------------------------------------------------------------------------- /data/flash_noflash/Shelves_009_flash.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Shelves_009_flash.png -------------------------------------------------------------------------------- /data/flash_noflash/Toys_007_ambient.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Toys_007_ambient.png -------------------------------------------------------------------------------- /data/flash_noflash/Toys_007_flash.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/Toys_007_flash.png -------------------------------------------------------------------------------- /data/flash_noflash/artifacts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/artifacts.png -------------------------------------------------------------------------------- /data/flash_noflash/result.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/flash_noflash/result.png -------------------------------------------------------------------------------- /data/gt_dirichlet.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/gt_dirichlet.png -------------------------------------------------------------------------------- /data/gt_neumann.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/gt_neumann.png -------------------------------------------------------------------------------- /data/laplace_sol_2d.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/laplace_sol_2d.png -------------------------------------------------------------------------------- /data/laplace_sol_3d.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/laplace_sol_3d.png -------------------------------------------------------------------------------- /data/poisson_image_editing/mask.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/poisson_image_editing/mask.jpg -------------------------------------------------------------------------------- /data/poisson_image_editing/result.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/poisson_image_editing/result.png -------------------------------------------------------------------------------- /data/poisson_image_editing/source.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/poisson_image_editing/source.jpg -------------------------------------------------------------------------------- /data/poisson_image_editing/target.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/poisson_image_editing/target.jpg -------------------------------------------------------------------------------- /data/sol_dirichlet.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/sol_dirichlet.png -------------------------------------------------------------------------------- /data/sol_neumann.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/sol_neumann.png -------------------------------------------------------------------------------- /data/zero_mean_false.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/zero_mean_false.png -------------------------------------------------------------------------------- /data/zero_mean_true.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/data/zero_mean_true.png -------------------------------------------------------------------------------- /examples/context.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 4 | 5 | import poissonpy -------------------------------------------------------------------------------- /examples/flash_noflash.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | from PIL import Image 4 | import matplotlib.pyplot as plt 5 | 6 | from context import poissonpy 7 | from poissonpy import functional, utils, solvers 8 | 9 | ambient_rgb = utils.read_image("../data/flash_noflash/Shelves_007_ambient.png", 0.25) 10 | flash_rgb = utils.read_image("../data/flash_noflash/Shelves_007_flash.png", 0.25) 11 | mask = np.ones_like(ambient_rgb[..., 0]) 12 | 13 | ambient_rgb = np.pad(ambient_rgb, ((1, 1), (1, 1), (0, 0)), constant_values=1) 14 | flash_rgb = np.pad(flash_rgb, ((1, 1), (1, 1), (0, 0)), constant_values=1) 15 | mask = np.pad(mask, ((1, 1), (1, 1))) 16 | 17 | res_rgb = [] 18 | for i in range(ambient_rgb.shape[-1]): 19 | ambient = ambient_rgb[..., i] 20 | flash = flash_rgb[..., i] 21 | 22 | # compute image gradients 23 | gx_a, gy_a = functional.get_np_gradient(ambient) 24 | gx_f, gy_f = functional.get_np_gradient(flash) 25 | 26 | # gradient coherence map 27 | M = np.abs(gx_a * gx_f + gy_a * gy_f) / (functional.get_np_gradient_amplitude(gx_a, gy_a) * functional.get_np_gradient_amplitude(gx_f, gy_f) + 1e-8) 28 | reversal = (gx_a * gx_f + gy_a * gy_f) / (functional.get_np_gradient_amplitude(gx_a, gy_a) * functional.get_np_gradient_amplitude(gx_f, gy_f) + 1e-8) < 0 29 | # saturation map 30 | w_s = utils.normalize(np.tanh(40 * (utils.normalize(flash) - 0.9))) 31 | 32 | t = np.abs(gx_a*gx_f + gy_a*gy_f) / (gx_f**2 + gy_f**2) 33 | gx_a_proj = t * gx_f 34 | gy_a_proj = t * gy_f 35 | gx_a_proj = np.where(reversal, -gx_a_proj, gx_a_proj) 36 | gy_a_proj = np.where(reversal, -gy_a_proj, gy_a_proj) 37 | #plt.imshow(gx_a_proj, cmap="gray") 38 | #plt.show() 39 | 40 | gx_a_new = gx_a_proj#w_s * gx_a + (1 - w_s) * (gx_a_proj) 41 | gy_a_new = gy_a_proj#w_s * gy_a + (1 - w_s) * (gy_a_proj) 42 | #gx_f_new = w_s * gx_a + (1 - w_s) * (M * gx_f + (1 - M) * gx_a) 43 | #gy_f_new = w_s * gy_a + (1 - w_s) * (M * gy_f + (1 - M) * gy_a) 44 | 45 | lap = functional.get_np_div(gx_a_new, gy_a_new) 46 | 47 | solver = solvers.Poisson2DRegion(mask, lap, ambient) 48 | 49 | res = solver.solve() 50 | res_rgb.append(res) 51 | 52 | res_rgb = np.dstack(res_rgb) 53 | res_rgb = np.clip(res_rgb, 0, 1) 54 | 55 | fig, axes = plt.subplots(1, 3) 56 | axes[0].imshow(ambient_rgb) 57 | axes[1].imshow(flash_rgb) 58 | axes[2].imshow(res_rgb) 59 | axes[0].set_title("ambient") 60 | axes[1].set_title("flash") 61 | axes[2].set_title("result") 62 | axes[0].axis("off") 63 | axes[1].axis("off") 64 | axes[2].axis("off") 65 | plt.tight_layout() 66 | plt.show() 67 | 68 | fig, axes = plt.subplots(1, 2) 69 | axes[0].imshow(ambient_rgb - res_rgb) 70 | axes[1].imshow(flash_rgb - res_rgb) 71 | axes[0].set_title("ambient artifacts") 72 | axes[1].set_title("flash artifacts") 73 | axes[0].axis("off") 74 | axes[1].axis("off") 75 | plt.tight_layout() 76 | plt.show() 77 | 78 | 79 | -------------------------------------------------------------------------------- /examples/numpy_mode.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/examples/numpy_mode.py -------------------------------------------------------------------------------- /examples/poisson_image_editing.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | from PIL import Image 4 | import matplotlib.pyplot as plt 5 | 6 | from context import poissonpy 7 | from poissonpy import functional, utils, solvers 8 | 9 | source = np.array(Image.open("../data/poisson_image_editing/source.jpg")).astype(np.float64) / 255 10 | mask = np.array(Image.open("../data/poisson_image_editing/mask.jpg").convert("L")).astype(np.float64) / 255 11 | target = np.array(Image.open("../data/poisson_image_editing/target.jpg")).astype(np.float64) / 255 12 | 13 | _, mask = cv2.threshold(mask, 0.5, 1, cv2.THRESH_BINARY) 14 | 15 | blended_rgb = [] 16 | # solve for each color channel 17 | for i in range(source.shape[-1]): 18 | # compute laplacian of interpolation function 19 | Gx_src, Gy_src = functional.get_np_gradient(source[..., i]) 20 | Gx_target, Gy_target = functional.get_np_gradient(target[..., i]) 21 | G_src_mag = (Gx_src**2 + Gy_src**2)**0.5 22 | G_target_mag = (Gx_target**2 + Gy_target**2)**0.5 23 | Gx = np.where(G_src_mag > G_target_mag, Gx_src, Gx_target) 24 | Gy = np.where(G_src_mag > G_target_mag, Gy_src, Gy_target) 25 | Gxx, _ = functional.get_np_gradient(Gx, forward=False) 26 | _, Gyy = functional.get_np_gradient(Gy, forward=False) 27 | laplacian = Gxx + Gyy 28 | 29 | # solve interpolation function 30 | solver = solvers.Poisson2DRegion(mask, laplacian, target[..., i]) 31 | solution = solver.solve() 32 | 33 | # alpha-blend interpolation and target function 34 | blended = mask * solution + (1 - mask) * target[..., i] 35 | blended_rgb.append(blended) 36 | 37 | blended_rgb = np.dstack(blended_rgb) 38 | blended_rgb = np.clip(blended_rgb, 0, 1) 39 | 40 | 41 | fig, axes = plt.subplots(1, 3) 42 | axes[0].imshow(source) 43 | axes[1].imshow(target) 44 | axes[2].imshow(blended_rgb) 45 | axes[0].set_title("source") 46 | axes[1].set_title("target") 47 | axes[2].set_title("result") 48 | axes[0].axis("off") 49 | axes[1].axis("off") 50 | axes[2].axis("off") 51 | plt.tight_layout() 52 | plt.show() -------------------------------------------------------------------------------- /examples/sympy_dirichlet.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from sympy import sin, cos 4 | from sympy.abc import x, y 5 | 6 | from context import poissonpy 7 | from poissonpy import functional, utils, solvers 8 | 9 | 10 | # analytic = sin(x) + cos(y) 11 | # laplacian = -sin(x) - cos(y) 12 | # x derivative = cos(x) 13 | # y derivative = -sin(y) 14 | f_expr = sin(x) + cos(y) 15 | laplacian_expr = functional.get_sp_laplacian_expr(f_expr) 16 | x_derivative_expr = functional.get_sp_derivative_expr(f_expr, x) 17 | y_derivative_expr = functional.get_sp_derivative_expr(f_expr, y) 18 | 19 | f = functional.get_sp_function(f_expr) 20 | laplacian = functional.get_sp_function(laplacian_expr) 21 | 22 | # possible boundary conditions: neumann_x, neumann_y, dirichlet 23 | interior = laplacian 24 | left = f 25 | right = f 26 | top = f 27 | bottom = f 28 | 29 | boundary = { 30 | "left": (left, "dirichlet"), 31 | "right": (right, "dirichlet"), 32 | "top": (top, "dirichlet"), 33 | "bottom": (bottom, "dirichlet") 34 | } 35 | 36 | solver = solvers.Poisson2DRectangle( 37 | ((-2*np.pi, -2*np.pi), (2*np.pi, 2*np.pi)), interior, boundary, 100, 100) 38 | gt = f(solver.x_grid, solver.y_grid) 39 | 40 | solution = solver.solve() 41 | err = np.abs(gt - solution) 42 | 43 | utils.plot_3d(solver.x_grid, solver.y_grid, gt, "ground truth") 44 | utils.plot_3d(solver.x_grid, solver.y_grid, solution, "solution") 45 | utils.plot_2d(solver.x, solver.y, solution, "solution") 46 | utils.plot_2d(solver.x, solver.y, err, "error") 47 | -------------------------------------------------------------------------------- /examples/sympy_laplace_eq.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from sympy import sin, cos 4 | from sympy.abc import x, y 5 | 6 | from context import poissonpy 7 | from poissonpy import functional, utils, solvers 8 | 9 | interior = 0 10 | left = functional.get_sp_function(sin(y)) 11 | right = functional.get_sp_function(sin(y)) 12 | top = functional.get_sp_function(sin(x)) 13 | bottom = functional.get_sp_function(sin(x)) 14 | 15 | boundary = { 16 | "left": (left, "dirichlet"), 17 | "right": (right, "dirichlet"), 18 | "top": (top, "dirichlet"), 19 | "bottom": (bottom, "dirichlet") 20 | } 21 | 22 | solver = solvers.Poisson2DRectangle( 23 | ((-2*np.pi, -2*np.pi), (2*np.pi, 2*np.pi)), interior, boundary, 100, 100) 24 | 25 | solution = solver.solve() 26 | 27 | utils.plot_2d(solver.x, solver.y, solution, "solution") 28 | utils.plot_3d(solver.x_grid, solver.y_grid, solution, "solution") 29 | utils.plot_3d(solver.x_grid, solver.y_grid, gt, "ground truth") 30 | 31 | -------------------------------------------------------------------------------- /examples/sympy_neumann.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from sympy import sin, cos 4 | from sympy.abc import x, y 5 | 6 | from context import poissonpy 7 | from poissonpy import functional, utils, solvers 8 | 9 | 10 | # analytic = sin(x) + cos(y) 11 | # laplacian = -sin(x) - cos(y) 12 | # x derivative = cos(x) 13 | # y derivative = -sin(y) 14 | f_expr = sin(x) + cos(y) 15 | laplacian_expr = functional.get_sp_laplacian_expr(f_expr) 16 | x_derivative_expr = functional.get_sp_derivative_expr(f_expr, x) 17 | y_derivative_expr = functional.get_sp_derivative_expr(f_expr, y) 18 | 19 | f = functional.get_sp_function(f_expr) 20 | laplacian = functional.get_sp_function(laplacian_expr) 21 | 22 | # possible boundary conditions: neumann_x, neumann_y, dirichlet 23 | interior = laplacian 24 | left = functional.get_sp_function(x_derivative_expr) 25 | right = f 26 | top = functional.get_sp_function(y_derivative_expr) 27 | bottom = f 28 | 29 | boundary = { 30 | "left": (left, "neumann_x"), 31 | "right": (right, "dirichlet"), 32 | "top": (top, "neumann_y"), 33 | "bottom": (bottom, "dirichlet") 34 | } 35 | 36 | solver = solvers.Poisson2DRectangle( 37 | ((-2*np.pi, -2*np.pi), (2*np.pi, 2*np.pi)), interior, boundary, 100, 100) 38 | gt = f(solver.x_grid, solver.y_grid) 39 | 40 | solution = solver.solve() 41 | err = np.abs(gt - solution) 42 | 43 | utils.plot_2d(solver.x, solver.y, solution, "solution") 44 | utils.plot_3d(solver.x_grid, solver.y_grid, solution, "solution") 45 | utils.plot_3d(solver.x_grid, solver.y_grid, gt, "ground truth") 46 | utils.plot_2d(solver.x, solver.y, err, "error") 47 | -------------------------------------------------------------------------------- /examples/sympy_region_laplace.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | 4 | from sympy import sin, cos 5 | from sympy.abc import x, y 6 | 7 | from context import poissonpy 8 | from poissonpy import functional, utils, solvers 9 | 10 | 11 | f_expr = sin(x**2 + y**2) / (x**2 + y**2) 12 | laplacian_expr = functional.get_sp_laplacian_expr(f_expr) 13 | 14 | f = functional.get_sp_function(f_expr) 15 | laplacian = functional.get_sp_function(laplacian_expr) 16 | 17 | interior = laplacian 18 | boundary = f 19 | mask = utils.circle_region(500, 500, 200) 20 | 21 | 22 | solver = solvers.Poisson2DRegion(mask ,interior, boundary, ((-5, -5), (5, 5))) 23 | gt = f(solver.x_grid, solver.y_grid) * mask 24 | 25 | solution = solver.solve() 26 | err = np.abs(gt - solution) 27 | 28 | utils.plot_3d(solver.x_grid, solver.y_grid, gt, "ground truth") 29 | utils.plot_3d(solver.x_grid, solver.y_grid, solution, "solution") 30 | utils.plot_2d(solver.x, solver.y, solution, "solution") 31 | utils.plot_2d(solver.x, solver.y, err, "error") -------------------------------------------------------------------------------- /examples/sympy_zero_mean.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from sympy import sin, cos 4 | from sympy.abc import x, y 5 | 6 | from context import poissonpy 7 | from poissonpy import functional, utils, solvers 8 | 9 | 10 | # analytic = sin(x) + cos(y) 11 | # laplacian = -sin(x) - cos(y) 12 | # x derivative = cos(x) 13 | # y derivative = -sin(y) 14 | f_expr = sin(x) + cos(y) 15 | laplacian_expr = functional.get_sp_laplacian_expr(f_expr) 16 | x_derivative_expr = functional.get_sp_derivative_expr(f_expr, x) 17 | y_derivative_expr = functional.get_sp_derivative_expr(f_expr, y) 18 | 19 | f = functional.get_sp_function(f_expr) 20 | laplacian = functional.get_sp_function(laplacian_expr) 21 | 22 | # possible boundary conditions: neumann_x, neumann_y, dirichlet 23 | interior = laplacian 24 | left = functional.get_sp_function(x_derivative_expr) 25 | right = functional.get_sp_function(x_derivative_expr) 26 | top = functional.get_sp_function(y_derivative_expr) 27 | bottom = functional.get_sp_function(y_derivative_expr) 28 | 29 | boundary = { 30 | "left": (left, "neumann_x"), 31 | "right": (right, "neumann_x"), 32 | "top": (top, "neumann_y"), 33 | "bottom": (bottom, "neumann_y") 34 | } 35 | 36 | solver = solvers.Poisson2DRectangle( 37 | ((-2*np.pi, -2*np.pi), (2*np.pi, 2*np.pi)), interior, boundary, 100, 100, zero_mean=True) 38 | gt = f(solver.x_grid, solver.y_grid) 39 | 40 | solution = solver.solve() 41 | err = np.abs(gt - solution) 42 | 43 | utils.plot_2d(solver.x, solver.y, solution, "solution") 44 | utils.plot_3d(solver.x_grid, solver.y_grid, solution, "solution") 45 | utils.plot_3d(solver.x_grid, solver.y_grid, gt, "ground truth") 46 | utils.plot_2d(solver.x, solver.y, err, "error") 47 | -------------------------------------------------------------------------------- /poissonpy/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/poissonpy/__init__.py -------------------------------------------------------------------------------- /poissonpy/functional.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import scipy.signal 4 | 5 | from sympy import lambdify, diff 6 | from sympy.abc import x, y 7 | 8 | # Function related 9 | def get_sp_function(expr): 10 | return lambdify([x, y], expr, "numpy") 11 | 12 | def get_sp_laplacian_expr(expr): 13 | return diff(expr, x, 2) + diff(expr, y, 2) 14 | 15 | def get_sp_derivative_expr(expr, var): 16 | return diff(expr, var, 1) 17 | 18 | def get_np_gradient(arr, dx=1, dy=1, forward=True): 19 | if forward: 20 | kx = np.array([ 21 | [0, 0, 0], 22 | [0, -1/dx, 1/dx], 23 | [0, 0, 0] 24 | ]) 25 | ky = np.array([ 26 | [0, 0, 0], 27 | [0, -1/dy, 0], 28 | [0, 1/dy, 0] 29 | ]) 30 | else: 31 | kx = np.array([ 32 | [0, 0, 0], 33 | [-1/dx, 1/dx, 0], 34 | [0, 0, 0] 35 | ]) 36 | ky = np.array([ 37 | [0, -1/dy, 0], 38 | [0, 1/dy, 0], 39 | [0, 0, 0] 40 | ]) 41 | Gx = scipy.signal.fftconvolve(arr, kx, mode="same") 42 | Gy = scipy.signal.fftconvolve(arr, ky, mode="same") 43 | return Gx, Gy 44 | 45 | def get_np_laplacian(arr, dx=1, dy=1): 46 | kernel = np.array([ 47 | [0, 1/(dy**2), 0], 48 | [1/(dx**2), -2/(dx**2)-2/(dy**2), 1/(dx**2)], 49 | [0, 1/(dy**2), 0] 50 | ]) 51 | laplacian = scipy.signal.fftconvolve(arr, kernel, mode="same") 52 | return laplacian 53 | 54 | def get_np_div(gx, gy, dx=1, dy=1): 55 | gxx, _ = get_np_gradient(gx, dx, dy, forward=False) 56 | _, gyy = get_np_gradient(gy, dx, dy, forward=False) 57 | return gxx + gyy 58 | 59 | def get_np_gradient_amplitude(gx, gy): 60 | return np.sqrt(gx**2 + gy**2) -------------------------------------------------------------------------------- /poissonpy/helpers.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from skimage.segmentation import find_boundaries 4 | 5 | def process_mask(mask): 6 | boundary = find_boundaries(mask, mode="inner").astype(int) 7 | inner = mask - boundary 8 | return inner, boundary 9 | 10 | def get_grid_ids(X, Y): 11 | grid_ids = np.arange(Y * X).reshape(Y, X) 12 | return grid_ids 13 | 14 | def get_selected_values(values, mask): 15 | assert values.shape == mask.shape 16 | nonzero_idx = np.nonzero(mask) # get mask 1 17 | return values[nonzero_idx] -------------------------------------------------------------------------------- /poissonpy/solvers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sympy as sp 3 | import scipy.sparse 4 | import types 5 | import matplotlib.pyplot as plt 6 | import seaborn as sns 7 | 8 | from . import functional, helpers 9 | 10 | # inner region conditions 11 | # laplacian diag matrix construction 12 | 13 | class Poisson2DRectangle: 14 | """ 15 | Solve 2D Poisson Equation on a rectangle 16 | """ 17 | def __init__(self, rect, interior, boundary, X=100, Y=100, zero_mean=False): 18 | """ 19 | Args: 20 | rect: 21 | interior: sympy expresson or numpy array 22 | boundary: 23 | N: Rectangle to N * N grid # 24 | """ 25 | self.x1, self.y1 = rect[0] # top-left corner 26 | self.x2, self.y2 = rect[1] # bottom-right corner 27 | self.X, self.Y = X, Y 28 | 29 | self.x = np.linspace(self.x1, self.x2, self.X) 30 | self.y = np.linspace(self.y1, self.y2, self.Y) 31 | 32 | self.dx = self.x[1] - self.x[0] 33 | self.dy = self.y[1] - self.y[0] 34 | 35 | self._x_grid, self._y_grid = np.meshgrid(self.x, self.y) 36 | self.xs = self.x_grid.flatten() 37 | self.ys = self.y_grid.flatten() 38 | 39 | self.interior = interior 40 | self.boundary = boundary 41 | for bd, (_, mode) in self.boundary.items(): 42 | assert bd in ["left", "right", "top", "bottom"] 43 | assert mode in ["dirichlet", "neumann_x", "neumann_y"] 44 | 45 | self.zero_mean = zero_mean 46 | 47 | self.A, self.b = self.build_linear_system() 48 | 49 | @property 50 | def x_grid(self): 51 | return self._x_grid 52 | 53 | @property 54 | def y_grid(self): 55 | return self._y_grid 56 | 57 | def build_linear_system(self): 58 | 59 | self.interior_ids = np.arange(self.X + 1, 2 * self.X - 1) + self.X * np.expand_dims(np.arange(self.Y - 2), 1) 60 | self.interior_ids = self.interior_ids.flatten() 61 | boundary_ids = { 62 | "left": self.X * np.arange(1, self.Y - 1), 63 | "right": self.X * np.arange(1, self.Y - 1) + (self.X - 1), 64 | "top": np.arange(1, self.X - 1), 65 | "bottom": np.arange(self.X * self.Y - self.X + 1, self.X * self.Y - 1) 66 | } 67 | 68 | self.all_ids = np.concatenate([self.interior_ids, np.concatenate(list(boundary_ids.values()))]) 69 | self.all_ids.sort() 70 | 71 | if self.zero_mean: 72 | A = scipy.sparse.lil_matrix((len(self.all_ids) + 1, len(self.all_ids) + 1)) 73 | b = np.zeros(len(self.all_ids) + 1) 74 | else: 75 | A = scipy.sparse.lil_matrix((len(self.all_ids), len(self.all_ids))) 76 | b = np.zeros(len(self.all_ids)) 77 | 78 | self.interior_pos = np.searchsorted(self.all_ids, self.interior_ids) 79 | boundary_pos = { 80 | bd: np.searchsorted(self.all_ids, boundary_ids[bd]) for bd in boundary_ids 81 | } 82 | 83 | # interior - laplacian 84 | n1_pos = np.searchsorted(self.all_ids, self.interior_ids - 1) 85 | n2_pos = np.searchsorted(self.all_ids, self.interior_ids + 1) 86 | n3_pos = np.searchsorted(self.all_ids, self.interior_ids - self.X) 87 | n4_pos = np.searchsorted(self.all_ids, self.interior_ids + self.X) 88 | 89 | # Discrete laplacian here important! 90 | A[self.interior_pos, n1_pos] = 1 / (self.dx**2) 91 | A[self.interior_pos, n2_pos] = 1 / (self.dx**2) 92 | A[self.interior_pos, n3_pos] = 1 / (self.dy**2) 93 | A[self.interior_pos, n4_pos] = 1 / (self.dy**2) 94 | A[self.interior_pos, self.interior_pos] = -2 / (self.dx**2) + -2 / (self.dy**2) 95 | 96 | if isinstance(self.interior, types.FunctionType): 97 | b[self.interior_pos] = self.interior(self.xs[self.interior_ids], self.ys[self.interior_ids]) 98 | elif isinstance(self.interior, (int, float)): 99 | b[self.interior_pos] = self.interior 100 | 101 | if self.zero_mean: 102 | A[-1, :] = 1.0 103 | A[:, -1] = 1.0 104 | A[-1, -1] = 0.0 105 | b[-1] = 0 106 | 107 | for bd, (bd_func, mode) in self.boundary.items(): 108 | bd_pos = boundary_pos[bd] 109 | bd_ids = boundary_ids[bd] 110 | 111 | if isinstance(bd_func, types.FunctionType): 112 | b[bd_pos] = bd_func(self.xs[bd_ids], self.ys[bd_ids]) 113 | elif isinstance(bd_func, (int, float)): 114 | b[bd_pos] = bd_func 115 | 116 | if mode == "dirichlet": 117 | A[bd_pos, bd_pos] = 1 118 | elif mode == "neumann_x": 119 | if bd == "left": 120 | n_ids = bd_ids + 1 121 | n_pos = np.searchsorted(self.all_ids, n_ids) 122 | A[bd_pos, bd_pos] = -1 / self.dx 123 | A[bd_pos, n_pos] = 1 / self.dx 124 | else: # right 125 | n_ids = bd_ids - 1 126 | n_pos = np.searchsorted(self.all_ids, n_ids) 127 | A[bd_pos, bd_pos] = 1 / self.dx 128 | A[bd_pos, n_pos] = -1 / self.dx 129 | elif mode == "neumann_y": 130 | if bd == "top": 131 | n_ids = bd_ids + self.X 132 | n_pos = np.searchsorted(self.all_ids, n_ids) 133 | A[bd_pos, bd_pos] = -1 / self.dy 134 | A[bd_pos, n_pos] = 1 / self.dy 135 | else: 136 | n_ids = bd_ids - self.X 137 | n_pos = np.searchsorted(self.all_ids, n_ids) 138 | A[bd_pos, bd_pos] = 1 / self.dy 139 | A[bd_pos, n_pos] = -1 / self.dy 140 | return A.tocsr(), b 141 | 142 | def solve(self): 143 | # multigrid solver result bad? 144 | #x = scipy.sparse.linalg.bicg(A, b)[0] 145 | x = scipy.sparse.linalg.spsolve(self.A, self.b) 146 | if self.zero_mean: 147 | x = x[:-1] 148 | 149 | grid_ids = np.arange(self.Y * self.X) 150 | all_pos = np.searchsorted(grid_ids, self.all_ids) 151 | 152 | solution_grid = np.zeros(self.Y * self.X) 153 | solution_grid[all_pos] = x 154 | solution_grid = solution_grid.reshape(self.Y, self.X) 155 | return solution_grid 156 | 157 | def estimated_laplacian(self, f): 158 | val = f(self.x_grid, self.y_grid).flatten() 159 | est_lap = (self.A @ val[self.all_ids])[self.interior_pos] 160 | grid_ids = np.arange(self.Y * self.X) 161 | interior_pos = np.searchsorted(grid_ids, self.interior_ids) 162 | solution_grid = np.zeros(self.Y * self.X) 163 | solution_grid[interior_pos] = est_lap 164 | solution_grid = solution_grid.reshape(self.Y, self.X) 165 | return solution_grid 166 | 167 | class Poisson2DRegion: 168 | """ 169 | Solve 2D Poisson Equation on a region with arbitrary shape. 170 | """ 171 | def __init__(self, region, interior, boundary, rect=None): 172 | self.region = region # region mask 173 | self.Y, self.X = self.region.shape 174 | 175 | if rect: 176 | x0, y0 = rect[0] 177 | x1, y1 = rect[1] 178 | self.x = np.linspace(x0, x1, self.X) 179 | self.y = np.linspace(y0, y1, self.Y) 180 | else: 181 | self.x = np.arange(self.X) 182 | self.y = np.arange(self.Y) 183 | self.x_grid, self.y_grid = np.meshgrid(self.x, self.y) 184 | 185 | self.dx = self.x[1] - self.x[0] 186 | self.dy = self.y[1] - self.y[0] 187 | 188 | if isinstance(interior, types.FunctionType): 189 | self.interior = interior(self.x_grid, self.y_grid) 190 | elif isinstance(interior, (int, float)): 191 | self.interior = np.ones_like(region) * interior 192 | else: 193 | self.interior = interior 194 | 195 | if isinstance(boundary, types.FunctionType): 196 | self.boundary = boundary(self.x_grid, self.y_grid) 197 | elif isinstance(boundary, (int, float)): 198 | self.boundary = np.ones_like(region) * boundary 199 | else: 200 | self.boundary = boundary 201 | 202 | self.A, self.b = self.build_linear_system() 203 | 204 | def build_linear_system(self): 205 | self.inner_region, self.boundary_region = helpers.process_mask(self.region) 206 | self.grid_ids = helpers.get_grid_ids(self.X, self.Y) 207 | 208 | 209 | self.inner_ids = helpers.get_selected_values(self.grid_ids, self.inner_region).flatten() 210 | self.boundary_ids = helpers.get_selected_values(self.grid_ids, self.boundary_region).flatten() 211 | self.region_ids = helpers.get_selected_values(self.grid_ids, self.region).flatten() # boundary + inner 212 | 213 | self.inner_pos = np.searchsorted(self.region_ids, self.inner_ids) 214 | self.boundary_pos = np.searchsorted(self.region_ids, self.boundary_ids) 215 | self.region_pos = np.searchsorted(self.grid_ids.flatten(), self.region_ids) 216 | 217 | n1_pos = np.searchsorted(self.region_ids, self.inner_ids - 1) 218 | n2_pos = np.searchsorted(self.region_ids, self.inner_ids + 1) 219 | n3_pos = np.searchsorted(self.region_ids, self.inner_ids - self.X) 220 | n4_pos = np.searchsorted(self.region_ids, self.inner_ids + self.X) 221 | 222 | A = scipy.sparse.lil_matrix((len(self.region_ids), len(self.region_ids))) 223 | A[self.inner_pos, n1_pos] = 1 / (self.dx**2) 224 | A[self.inner_pos, n2_pos] = 1 / (self.dx**2) 225 | A[self.inner_pos, n3_pos] = 1 / (self.dy**2) 226 | A[self.inner_pos, n4_pos] = 1 / (self.dy**2) 227 | A[self.inner_pos, self.inner_pos] = -2 / (self.dx**2) + -2 / (self.dy**2) 228 | 229 | A[self.boundary_pos, self.boundary_pos] = 1 # only dirichlet for now 230 | A = A.tocsr() 231 | 232 | boundary_conditions = helpers.get_selected_values(self.boundary, self.boundary_region).flatten() 233 | interior_laplacians = helpers.get_selected_values(self.interior, self.inner_region).flatten() 234 | b = np.zeros(len(self.region_ids)) 235 | b[self.inner_pos] = interior_laplacians 236 | b[self.boundary_pos] = boundary_conditions 237 | 238 | return A, b 239 | 240 | def solve(self): 241 | # multigrid solver result bad? 242 | #x = scipy.sparse.linalg.bicg(A, b)[0] 243 | x = scipy.sparse.linalg.spsolve(self.A, self.b) 244 | 245 | solution_grid = np.zeros(self.Y * self.X) 246 | solution_grid[self.region_pos] = x 247 | solution_grid = solution_grid.reshape(self.Y, self.X) 248 | return solution_grid 249 | -------------------------------------------------------------------------------- /poissonpy/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | from PIL import Image 3 | import numpy as np 4 | import seaborn as sns 5 | import matplotlib.pyplot as plt 6 | from matplotlib import cm 7 | from matplotlib.ticker import LinearLocator 8 | 9 | # Image utils 10 | def read_image(filename, scale=1): 11 | img = Image.open(os.path.join(filename)) 12 | if scale != 1: 13 | w, h = img.size 14 | img = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC) 15 | img = np.array(img) 16 | if len(img.shape) == 3: 17 | img = img[..., :3] 18 | return img.astype(np.float64) / 255 # only first 3 19 | 20 | def normalize(img): 21 | return (img - img.min()) / (img.max() - img.min()) 22 | 23 | def circle_region(X, Y, r): 24 | o_x, o_y = X // 2, Y // 2 25 | x_grid, y_grid = np.meshgrid(np.arange(X), np.arange(Y)) 26 | circ = (((x_grid - o_x)**2 + (y_grid - o_y)**2) < r**2).astype(np.int) 27 | return circ 28 | 29 | # Plotting 30 | def plot_2d(X, Y, Z, title): 31 | fig, ax = plt.subplots() 32 | heatmap = ax.imshow(Z[::-1], cmap=cm.coolwarm) 33 | ax.set_title(title) 34 | ax.set_xlabel("x") 35 | ax.set_ylabel("y") 36 | ax.set_xticks([0, len(X) - 1], np.round(np.array([X[0], X[-1]]), 2)) 37 | ax.set_yticks([0, len(Y) - 1], np.round(np.array([Y[-1], Y[0]]), 2)) 38 | fig.colorbar(heatmap, shrink=0.5, aspect=5) 39 | plt.show() 40 | 41 | def plot_3d(X, Y, Z, title): 42 | fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) 43 | 44 | # Plot the surface. 45 | surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm, 46 | linewidth=0, antialiased=False) 47 | 48 | # Customize the z axis. 49 | ax.zaxis.set_major_locator(LinearLocator(10)) 50 | # A StrMethodFormatter is used automatically 51 | ax.zaxis.set_major_formatter('{x:.02f}') 52 | 53 | # Add a color bar which maps values to colors. 54 | fig.colorbar(surf, shrink=0.5, aspect=5) 55 | 56 | plt.title(title) 57 | plt.show() 58 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bchao1/poissonpy/81b1c30c30e13052bc18f3a4e693a2295301e213/requirements.txt -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup( 4 | name='poissonpy', 5 | version='1.0.0', 6 | description='Plug-and-play 2D Poisson equations library.', 7 | long_description="Plug-and-play 2D Poisson equation library. Useful in scientific computing, image and video processing, computer graphics.", 8 | url='https://github.com/bchao1/poissonpy', 9 | keywords = "pde poisson-equation laplace-equation scientific-computing poisson", 10 | author='bchao1', 11 | license='MIT', 12 | packages=find_packages(), 13 | python_requires=">=3.6" 14 | ) --------------------------------------------------------------------------------