CODEWITHSUNDEEP
pythonnumpymatplotlib
aignas
aignas

Reputation: 1074

How to "draw" a geometric shape on a 2-D meshgrid?

I want to do a simulation on Fraunhofer diffraction and for that purpose I chose to use numpy and matplotlib. What I need to do, is to specify a 2D aperture function and for that I can create a meshgrid of x and y values and assign a function z(x,y), which in this case should be complex. All this does not sound too complicated, but here's where I bump into a problem.

How do you define a rectangular, or a triangular piece inside a meshgrid, such that inside the geometric figure z=1 and outside z=0?

Minimal working example, where to begin:

#! /bin/usr/env python

# Import environment
import numpy as np

x_ = np.linspace(0,1,255)
y_ = np.linspace(0,1,255)
x,y = np.meshgrid(x_,y_)

what to do next?

I tried to solve the problem differently:

However, this puts severe restrictions on the values taken by the function z, which is the main reason I am searching for a different approach.

Thanks very much for anybody who can help me.

Upvotes: 3

Views: 2949

Answers (1)

Joe Kington
Joe Kington

Reputation: 284552

The two easiest options are to either use matplotlib.nxutils.points_inside_poly or to use mahotas.polygon.fill_polygon. The latter is a bit faster, but requires installing mahotas.

As an example of the first option:

import numpy as np
from matplotlib.nxutils import points_inside_poly

nx, ny = 10, 10
poly_verts = [(1,1), (5,1), (5,9),(3,2),(1,1)]

# Create vertex coordinates for each grid cell...
# (<0,0> is at the top left of the grid in this system)
x, y = np.meshgrid(np.arange(nx), np.arange(ny))
x, y = x.flatten(), y.flatten()

points = np.vstack((x,y)).T

grid = points_inside_poly(points, poly_verts)
grid = grid.reshape((ny,nx))

print grid

Which yields (a boolean numpy array):

[[False False False False False False False False False False]
 [False  True  True  True  True False False False False False]
 [False False False  True  True False False False False False]
 [False False False False  True False False False False False]
 [False False False False  True False False False False False]
 [False False False False  True False False False False False]
 [False False False False False False False False False False]
 [False False False False False False False False False False]
 [False False False False False False False False False False]
 [False False False False False False False False False False]]

On a side note, nxutils is going to be depreciated at some point in favor of some of the path methods. In the future, you'll probably want to do something along the lines of:

from matplotlib import path

...
p = path.Path(poly_verts)
grid = p.contains_points(points)
...

However, that's only in the github head at the moment.

Upvotes: 2

Related Questions