Add origin to Zernike

CHANGELOG.rst

@@ -10,6 +10,7 @@ New Features
 ------------
 - Added intersect/reflect/refract methods directly to `Optic.Interface`
   objects.
+- Added x_origin and y_origin arguments to batoid.Zernike.
 
 Performance Improvements
 ------------------------

batoid/surface.py

@@ -453,17 +453,21 @@ class Zernike(Surface):
     ----------
     coef : list of float
         Annular Zernike polynomial coefficients.
-    R_outer : float
+    R_outer : float, optional
         Outer radius of annulus.
-    R_inner : float
+    R_inner : float, optional
         Inner radius of annulus.
+    x_origin, y_origin : float, optional
+        Origin of the Zernike polynomial.  Default is (0, 0).
     """
-    def __init__(self, coef, R_outer=1.0, R_inner=0.0):
+    def __init__(self, coef, R_outer=1.0, R_inner=0.0, x_origin=0.0, y_origin=0.0):
         import galsim
 
         self.coef = np.array(coef, dtype=float, order="C")
         self.R_outer = float(R_outer)
         self.R_inner = float(R_inner)
+        self.x_origin = float(x_origin)
+        self.y_origin = float(y_origin)
         self.Z = galsim.zernike.Zernike(coef, R_outer, R_inner)
         self._xycoef = self.Z._coef_array_xy
         self._xycoef_gradx = self.Z.gradX._coef_array_xy
@@ -473,30 +477,45 @@ def __init__(self, coef, R_outer=1.0, R_inner=0.0):
             self._xycoef.ctypes.data,
             self._xycoef_gradx.ctypes.data,
             self._xycoef_grady.ctypes.data,
+            self.x_origin, self.y_origin,
             self._xycoef.shape[0],
             self._xycoef.shape[1]
         )
 
     def __hash__(self):
         return hash((
             "batoid.Zernike",
-            tuple(self.coef), self.R_outer, self.R_inner
+            tuple(self.coef),
+            self.R_outer, self.R_inner,
+            self.x_origin, self.y_origin,
         ))
 
     def __setstate__(self, args):
         self.__init__(*args)
 
     def __getstate__(self):
-        return self.coef, self.R_outer, self.R_inner
+        return self.coef, self.R_outer, self.R_inner, self.x_origin, self.y_origin
 
     def __eq__(self, rhs):
         if not isinstance(rhs, Zernike): return False
         return (np.array_equal(self.coef, rhs.coef) and
                 self.R_outer == rhs.R_outer and
-                self.R_inner == rhs.R_inner)
+                self.R_inner == rhs.R_inner and
+                self.x_origin == rhs.x_origin and
+                self.y_origin == rhs.y_origin)
 
     def __repr__(self):
-        return f"Zernike({self.coef!r}, {self.R_outer}, {self.R_inner})"
+        out = f"Zernike({self.coef!r}"
+        if self.R_outer != 1.0:
+            out += f", R_outer={self.R_outer}"
+        if self.R_inner != 0.0:
+            out += f", R_inner={self.R_inner}"
+        if self.x_origin != 0.0:
+            out += f", x_origin={self.x_origin}"
+        if self.y_origin != 0.0:
+            out += f", y_origin={self.y_origin}"
+        out += ")"
+        return out
 
 
 class Bicubic(Surface):

include/polynomialSurface.h

@@ -13,6 +13,7 @@ namespace batoid {
     public:
         PolynomialSurface(
             const double* coefs, const double* coefs_gradx, const double* coefs_grady,
+            const double xorigin, const double yorigin,
             size_t xsize, size_t ysize
         );
         ~PolynomialSurface();
@@ -29,6 +30,7 @@ namespace batoid {
         const double* _coefs;
         const double* _coefs_gradx;
         const double* _coefs_grady;
+        const double _xorigin, _yorigin;
         const size_t _xsize, _ysize;
     };
 

pysrc/polynomialSurface.cpp

@@ -12,13 +12,16 @@ namespace batoid {
                     size_t coefs,
                     size_t coefs_gradx,
                     size_t coefs_grady,
+                    const double xorigin,
+                    const double yorigin,
                     size_t xsize,
                     size_t ysize
                 ){
                     return new PolynomialSurface(
                         reinterpret_cast<const double*>(coefs),
                         reinterpret_cast<const double*>(coefs_gradx),
                         reinterpret_cast<const double*>(coefs_grady),
+                        xorigin, yorigin,
                         xsize, ysize
                     );
                 }

src/polynomialSurface.cpp

@@ -9,10 +9,12 @@ namespace batoid {
 
     PolynomialSurface::PolynomialSurface(
         const double* coefs, const double* coefs_gradx, const double* coefs_grady,
+        const double xorigin, const double yorigin,
         size_t xsize, size_t ysize
     ) :
         Surface(),
         _coefs(coefs), _coefs_gradx(coefs_gradx), _coefs_grady(coefs_grady),
+        _xorigin(xorigin), _yorigin(yorigin),
         _xsize(xsize), _ysize(ysize)
     {}
 
@@ -31,15 +33,15 @@ namespace batoid {
     }
 
     double PolynomialSurface::sag(double x, double y) const {
-        return horner2d(x, y, _coefs, _xsize, _ysize);
+        return horner2d(x-_xorigin, y-_yorigin, _coefs, _xsize, _ysize);
     }
 
     void PolynomialSurface::normal(double x, double y, double& nx, double& ny, double& nz) const {
         // The gradient arrays are always shaped 1 row and column smaller than the coef array.
         // The only exception is when the coef array is 1x1, in which case the gradient array is
         // also 1x1, but filled with a zero (so horner still returns the right result).
-        nx = -horner2d(x, y, _coefs_gradx, _xsize-1, _ysize-1);
-        ny = -horner2d(x, y, _coefs_grady, _xsize-1, _ysize-1);
+        nx = -horner2d(x-_xorigin, y-_yorigin, _coefs_gradx, _xsize-1, _ysize-1);
+        ny = -horner2d(x-_xorigin, y-_yorigin, _coefs_grady, _xsize-1, _ysize-1);
         nz = 1./std::sqrt(nx*nx + ny*ny + 1);
         nx *= nz;
         ny *= nz;
@@ -81,7 +83,9 @@ namespace batoid {
                     map(to:coefs[:size], coefs_gradx[:sizem1], coefs_grady[:sizem1])
                 #pragma omp target map(from:ptr)
                 {
-                    ptr = new PolynomialSurface(coefs, coefs_gradx, coefs_grady, _xsize, _ysize);
+                    ptr = new PolynomialSurface(
+                        coefs, coefs_gradx, coefs_grady, _xorigin, _yorigin, _xsize, _ysize
+                    );
                 }
                 _devPtr = ptr;
             }

tests/test_Zernike.py

@@ -79,6 +79,15 @@ def test_sag():
             rtol=0,
             atol=1e-12
         )
+        # Check origin
+        bz2 = batoid.Zernike(coef, R_outer=R_outer, R_inner=R_inner, x_origin=0.1, y_origin=0.2)
+        do_pickle(bz2)
+        np.testing.assert_allclose(
+            bz2.sag(x, y),
+            bz.sag(x-0.1, y-0.2),
+            rtol=0,
+            atol=1e-12
+        )
 
 
 @timer