mperrin / poppy

Physical Optics Propagation in Python
BSD 3-Clause "New" or "Revised" License
177 stars 41 forks source link

code submission: generate orthonormal basis on arbitrary aperture #166

Closed mperrin closed 8 years ago

mperrin commented 8 years ago

From Arthur Vigan at LAM:

Dear Marshall,

I have started using your poppy library, mainly for the implementation of the Zernike polynomials that you have included in the library. In particular, you have written an orthonormalisation process to obtain a basis on an hexagonal aperture, which is super useful! But I think it would be even better to have something for an arbitrary aperture. I recently needed to have Zernikes on a circular aperture with central obscuration, so I implemented a new arbitrary_basis() method (see code below). Please feel free if you think it's a useful addition to the library.


####################################################################################

def arbitrary_basis(aperture, nterms=15, rho=None, theta=None):
   """
   Return a cube of Zernike terms from 1 to N, calculated on an
   arbitrary aperture, each as a 2D array showing the value at each
   point. (Regions outside the unit circle on which the Zernike is
   defined are initialized to zero.)

   Parameters
   -----------
   aperture : array_like
       2D binary array representing the arbitrary aperture
   nterms : int, optional
       Number of Zernike terms to return, starting from piston.
       (e.g. ``nterms=1`` would return only the Zernike piston term.)
       Default is 15.
   rho, theta : array_like, optional
       Image plane coordinates. `rho` should be 0 at the origin
       and 1.0 at the edge of the pupil. `theta` should be
       the angle in radians.
   """

   shape = aperture.shape
   npix  = shape[0]

   A = aperture.sum()

   # precompute zernikes
   Z = np.zeros((nterms + 1,) + shape)
   Z[1:] = zernike_basis(nterms=nterms, npix=npix, rho=rho, theta=theta, outside=0.0)

   G = [np.zeros(shape), np.ones(shape)]  # array of G_i etc. intermediate fn
   H = [np.zeros(shape), np.ones(shape) * aperture]  # array of hexikes
   c = {}  # coefficients hash

   for j in np.arange(nterms - 1) + 1:  # can do one less since we already have the piston term
       _log.debug("  j = " + str(j))
       # Compute the j'th G, then H
       nextG = Z[j + 1] * aperture
       for k in np.arange(j) + 1:
           c[(j + 1, k)] = -1 / A * (Z[j + 1] * H[k] * aperture).sum()
           if c[(j + 1, k)] != 0:
               nextG += c[(j + 1, k)] * H[k]
           _log.debug("    c[%s] = %f", str((j + 1, k)), c[(j + 1, k)])

       nextH = nextG / np.sqrt((nextG ** 2).sum() / A)

       G.append(nextG)
       H.append(nextH)

       #TODO - contemplate whether the above algorithm is numerically stable
       # cf. modified gram-schmidt algorithm discussion on wikipedia.

   # drop the 0th null element, return the rest
   return H[1:]
mperrin commented 8 years ago

added; not really tested particularly.