Various cleanup in Optical class

MANIFEST.in

@@ -1,6 +1,7 @@
 include *.rst
 include LICENSE
 recursive-include tests *.py
+recursive-include share *
 prune *docs/_build
 prune *tests/input
 prune *tests/output

piff/__init__.py

@@ -114,3 +114,5 @@
 # Leave these in their own namespaces
 from . import util
 from . import des
+from . import wavefront
+from . import meta_data

piff/meta_data.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2016 by Mike Jarvis and the other collaborators on GitHub at
+# https://github.com/rmjarvis/Piff  All rights reserved.
+#
+# Piff is free software: Redistribution and use in source and binary forms
+# with or without modification, are permitted provided that the following
+# conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the disclaimer given in the accompanying LICENSE
+#    file.
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the disclaimer given in the documentation
+#    and/or other materials provided with the distribution.
+
+import os
+
+piff_dir = os.path.split(os.path.realpath(__file__))[0]
+
+if 'PIFF_DATA_DIR' in os.environ:
+    data_dir = os.environ['PIFF_DATA_DIR']
+else:
+    data_dir = os.path.join(piff_dir, 'share')

piff/share

@@ -0,0 +1 @@
+../share/

piff/wavefront.py

@@ -5,8 +5,9 @@
 
 def convert_zernikes_des(a_fp):
     """ This method converts an array of Noll Zernike coefficients from
-    Focal Plane coordinates to a set suitable for u,v Sky coordinates. Assumes
-    that the WCS is of form [[0,-1],[-1,0]].  See AR's notebook calculate-Donut-to-GalsimPiff-Zernike-Coordinate-conversion
+    Focal Plane coordinates to a set suitable for u,v Sky coordinates.
+    Assumes that the WCS is of form [[0,-1],[-1,0]].
+    See AR's notebook calculate-Donut-to-GalsimPiff-Zernike-Coordinate-conversion
     for the calculation of this conversion.
 
     param: a_fp       An ndarray or list with Zernike's in Focal Plane coordinates
@@ -71,9 +72,10 @@ def __init__(self,wavefront_kwargs,logger=None):
         """ Parse the input options
 
         param: wavefront_kwargs    A dictionary holding the options for each
-                                   source of Zernike Coefficients.  Multiple input files are allowed,
-                                   with Dictionaries keyed by 'source1','source2'...
-                                   Each 'sourceN' dictionary has keys: 'file','zlist','keys','chip','wavelength'
+                                   source of Zernike Coefficients.  Multiple input files are
+                                   allowed, with Dictionaries keyed by 'source1','source2'...
+                                   Each 'sourceN' dictionary has keys:
+                                        'file','zlist','keys','chip','wavelength'
                                    The key 'survey' applies custom code for the desired survey.
         param: logger              A logger object for logging debug info. [default: None]
         """
@@ -148,8 +150,9 @@ def __init__(self,wavefront_kwargs,logger=None):
                         Z = zarray[achip,:,:,iZ]
 
                         # make interpolation object for each desired Zernike coefficient
-                        self.interp_objects[(isource,achip,iZ)] = LookupTable2D(xarray[achip,:],yarray[achip,:],Z,interpolant='spline')
-
+                        tab = LookupTable2D(xarray[achip,:], yarray[achip,:], Z,
+                                            interpolant='spline')
+                        self.interp_objects[(isource,achip,iZ)] = tab
             else:
 
                 # store None in chipkeys
@@ -161,17 +164,19 @@ def __init__(self,wavefront_kwargs,logger=None):
                     Z = zarray[:,:,iZ]
 
                     # make interpolation object for each desired Zernike coefficient
-                    self.interp_objects[(isource,None,iZ)] = LookupTable2D(xarray,yarray,Z,interpolant='spline')
+                    tab = LookupTable2D(xarray, yarray, Z, interpolant='spline')
+                    self.interp_objects[(isource,None,iZ)] = tab
 
-        return
 
     def fillWavefront(self,star_list,wavelength=-1.0,logger=None,addtostars=True):
-        """ Interpolate wavefront to each star's location, fill wavefront key of star.data.properties
-
-        :param star_list:         A list of Star instances
-        :param wavelength:        A float with the wavelength in [nm] to rescale the wavefront. [default: -1.0]
-        :param logger:            A logger object for logging debug info. [default: None]
-        :param addtostars:        A boolean flag to return stars with wavefront added to star.data or return array of wavefronts. [default: True]
+        """Interpolate wavefront to each star's location, fill wavefront key of star.data.properties
+
+        :param star_list:       A list of Star instances
+        :param wavelength:      A float with the wavelength in [nm] to rescale the wavefront.
+                                [default: -1.0]
+        :param logger:          A logger object for logging debug info. [default: None]
+        :param addtostars:      A boolean flag to return stars with wavefront added to star.data or
+                                return array of wavefronts. [default: True]
         """
 
         logger = galsim.config.LoggerWrapper(logger)

setup.py

@@ -49,11 +49,7 @@ def find_packages(base_path='.'):
 
 scripts = ['piffify', 'plotify', 'meanify']
 scripts = [ os.path.join('scripts',f) for f in scripts ]
-
-sources = glob.glob(os.path.join('src','*.c'))
-sources += glob.glob(os.path.join('src','*.cpp'))
-
-headers = glob.glob(os.path.join('include','*.h'))
+shared_data = glob.glob('share/*')
 
 undef_macros = []
 
@@ -278,11 +274,6 @@ def run(self):
         install_scripts.run(self)
         self.distribution.script_install_dir = self.install_dir
 
-ext=Extension("piff._piff",
-              sources,
-              depends=headers,
-              undef_macros = undef_macros)
-
 dependencies = ['galsim>=2.3', 'numpy>=1.17', 'scipy>=1.2', 'pyyaml>=5.1', 'treecorr>=4.3.1', 'fitsio>=1.0', 'matplotlib>=3.3', 'LSSTDESC.Coord>=1.0', 'treegp>=0.6', 'threadpoolctl>=3.1']
 
 with open('README.rst') as file:
@@ -310,16 +301,12 @@ def run(self):
       url="https://github.com/rmjarvis/Piff",
       download_url="https://github.com/rmjarvis/Piff/releases/tag/v%s.zip"%piff_version,
       packages=packages,
+      package_data={'piff' : shared_data},
       install_requires=dependencies,
-      # The rest of these are used by TreeCorr.  We might at some point want to do something
-      # like this here.  But for now, just comment them out.
-      #data_files=[('piff/include',headers)],
-      #ext_modules=[ext],
       cmdclass = {'build_ext': my_builder,
                   'install_scripts': my_install_scripts,
                   'easy_install': my_easy_install,
                   },
-      #headers=headers,
       scripts=scripts
       )
 

tests/test_gsobject_model.py

@@ -18,6 +18,7 @@
 import numpy as np
 import os
 import fitsio
+from unittest import mock
 
 from piff_test_helper import timer, CaptureLog
 
@@ -949,7 +950,6 @@ def test_fail():
 
     # reflux is harder to make fail.  Rather than try something even more contrived,
     # just mock np.linalg.solve to simulate the case that AtA ends up singular.
-    from unittest import mock
     with mock.patch('numpy.linalg.solve', side_effect=np.linalg.LinAlgError) as raising_solve:
         with CaptureLog(3) as cl:
             stars, nremoved = psf.reflux_stars([star1], logger=cl.logger)

tests/test_optics.py

@@ -19,6 +19,8 @@
 import os
 import fitsio
 import galsim
+import importlib
+from unittest import mock
 
 from piff_test_helper import timer
 
@@ -106,34 +108,39 @@ def test_optical(model=None):
     # test Zernike kwargs
     zernike_coeff_short = [0.,0.,0.,0.,1.,1.,1.,1.,1.,1.,1.,1.]
     nz = len(zernike_coeff_short)
-    param_test = model.kwargs_to_params(zernike_coeff=zernike_coeff_short,r0=0.12,g1=-0.05,g2=0.03,L0=20.)
+    param_test = model.kwargs_to_params(zernike_coeff=zernike_coeff_short,
+                                        r0=0.12, g1=-0.05, g2=0.03, L0=20.)
     for i in range(nz):
         assert zernike_coeff_short[i]==param_test[model.idx_z0+i]
     for i in range(nz+1,37+1):
         assert param_test[model.idx_z0+i]==0.0
 
     zernike_coeff_long = [0.,0.,0.,0.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,10.]
     nz = len(zernike_coeff_long)
-    param_test = model.kwargs_to_params(zernike_coeff=zernike_coeff_long,r0=0.12,g1=-0.05,g2=0.03,L0=20.)
+    param_test = model.kwargs_to_params(zernike_coeff=zernike_coeff_long,
+                                        r0=0.12, g1=-0.05, g2=0.03, L0=20.)
     for i in range(37+1):
         assert zernike_coeff_long[i]==param_test[model.idx_z0+i]
     assert param_test[model.idx_r0]==0.12
 
     # Test gsparams
-    gsp_star = galsim.GSParams(minimum_fft_size=32, folding_threshold=0.02)
-    gsp_donut = galsim.GSParams(minimum_fft_size=128, folding_threshold=0.005)
-    model = piff.Optical(template='des', gsparams=gsp_star)
-    assert model.gsparams == gsp_star
-    model = piff.Optical(template='des', atmo_type='Kolmogorov', gsparams='star')
-    assert model.gsparams == gsp_star
-    model = piff.Optical(template='des_donut', gsparams='donut')
-    assert model.gsparams == gsp_donut
-    model = piff.Optical(template='des_donut', minimum_fft_size=128, folding_threshold=0.005)
-    assert model.gsparams == gsp_donut
+    gsp_default = galsim.GSParams(minimum_fft_size=128, folding_threshold=0.005)
+    gsp_fast = galsim.GSParams(minimum_fft_size=64, folding_threshold=0.01)
+    gsp_faster = galsim.GSParams(minimum_fft_size=32, folding_threshold=0.02)
+    model = piff.Optical(template='des', gsparams=gsp_fast)
+    assert model.gsparams == gsp_fast
+    model = piff.Optical(template='des', atmo_type='Kolmogorov', gsparams='faster')
+    assert model.gsparams == gsp_faster
+    model = piff.Optical(template='des_donut')
+    assert model.gsparams is None
+    model = piff.Optical(template='des_donut', gsparams=galsim.GSParams())
+    assert model.gsparams == gsp_default
+    model = piff.Optical(template='des_donut', minimum_fft_size=64, folding_threshold=0.01)
+    assert model.gsparams == gsp_fast
     with np.testing.assert_raises(ValueError):
         piff.Optical(gsparams='invalid')
     with np.testing.assert_raises(ValueError):
-        piff.Optical(gsparams='donut', minimum_fft_size=128)
+        piff.Optical(gsparams='fast', minimum_fft_size=128)
 
 @timer
 def test_draw():
@@ -156,36 +163,103 @@ def test_draw():
     astar = stars[10]
 
     # test draw
-    astar.fit.params = model.kwargs_to_params(zernike_coeff=[0.,0.,0.,0.,0.2],r0=0.12,L0=10.,g1=0.01,g2=-0.02)
+    astar.fit.params = model.kwargs_to_params(zernike_coeff=[0.,0.,0.,0.,0.2],
+                                              r0=0.12, L0=10., g1=0.01, g2=-0.02)
     astar1 = model.draw(astar)
 
 @timer
-def test_pupil_im(pupil_plane_file='input/DECam_pupil_512uv.fits'):
-    import galsim
+def test_pupil_im(pupil_plane_file='DECam_pupil_512uv.fits'):
     print('test pupil im: ', pupil_plane_file)
     # make sure we can load up a pupil image
-    model = piff.Optical(diam=4.010, lam=500., pupil_plane_im=pupil_plane_file,atmo_type='Kolmogorov')
+    model = piff.Optical(diam=4.010, lam=500., pupil_plane_im=pupil_plane_file,
+                         atmo_type='Kolmogorov')
     test_optical(model)
     # make sure we really loaded it
-    pupil_plane_im = galsim.fits.read(pupil_plane_file)
+    # Note: piff automatically finds this file in the installed share directory.
+    full_pupil_plane_file = os.path.join('../share', pupil_plane_file)
+    pupil_plane_im = galsim.fits.read(full_pupil_plane_file)
     # Check the scale (and fix if necessary)
     print('pupil_plane_im.scale = ',pupil_plane_im.scale)
-    ref_psf = galsim.OpticalPSF(lam=500., diam=4.020)
-    print('scale should be ',ref_psf._psf.aper.pupil_plane_scale)
-    if pupil_plane_im.scale != ref_psf._psf.aper.pupil_plane_scale:
+    ref_aper = galsim.Aperture(diam=4.010, lam=500, pupil_plane_im=pupil_plane_im.array)
+    ref_aper._load_pupil_plane()
+    print('scale should be ',ref_aper.pupil_plane_scale)
+    if pupil_plane_im.scale != ref_aper.pupil_plane_scale:
         print('fixing scale')
-        pupil_plane_im.scale = ref_psf._psf.aper.pupil_plane_scale
-        pupil_plane_im.write(pupil_plane_file)
+        pupil_plane_im.scale = ref_aper.pupil_plane_scale
+        pupil_plane_im.write(full_pupil_plane_file)
 
-    model_pupil_plane_im = model.opt_kwargs['pupil_plane_im']
-    np.testing.assert_array_equal(pupil_plane_im.array, model_pupil_plane_im)
+    np.testing.assert_array_equal(pupil_plane_im.array, model.aperture._pupil_plane_im.array)
+
+    # Can also give a the full path, rather than automatically find it in the piff data_dir
+    model = piff.Optical(diam=4.010, lam=500.,
+                         pupil_plane_im=full_pupil_plane_file,
+                         atmo_type='Kolmogorov')
+    model.aperture._load_pupil_plane()
+    assert pupil_plane_im == model.aperture._pupil_plane_im
 
-    # test passing a different optical template that includes diam
-    piff.optical_model.optical_templates['test'] = {'diam': 2, 'lam':500}
-    model = piff.Optical(pupil_plane_im=pupil_plane_im, template='test')
+    print('Try diam=2 model')
+    model = piff.Optical(pupil_plane_im=pupil_plane_im, diam=2, lam=500)
     model_pupil_plane_im = model.opt_kwargs['pupil_plane_im']
-    np.testing.assert_array_equal(pupil_plane_im.array, model_pupil_plane_im)
+    np.testing.assert_array_equal(pupil_plane_im.array, model_pupil_plane_im.array)
+
+    # Error if file not found.
+    with np.testing.assert_raises(ValueError):
+        model = piff.Optical(diam=4.010, lam=500., pupil_plane_im='invalid.fits')
+
+    with mock.patch('os.environ', {'PIFF_DATA_DIR': 'invalid'}):
+        importlib.reload(piff.meta_data)
+        assert piff.meta_data.data_dir == 'invalid'
+        with np.testing.assert_raises(ValueError):
+            model = piff.Optical(diam=4.010, lam=500., pupil_plane_im=pupil_plane_file)
+    importlib.reload(piff.meta_data)
+
+
+@timer
+def test_mirror_figure():
+    mirror_figure_file = 'DECam_236392_finegrid512_nm_uv.fits'
+    full_mirror_figure_file = os.path.join('../share', mirror_figure_file)
+    mirror_figure_im = galsim.fits.read(full_mirror_figure_file)
+
+    # Check scale and fix if necessary.
+    mirror_figure_halfsize = 2.22246
+    ref_scale = 2 * mirror_figure_halfsize / 511
+    print('mirror_figure_scale should be ',ref_scale)
+    if mirror_figure_im.scale != ref_scale:
+        print('fixing scale')
+        mirror_figure_im.scale = ref_scale
+        mirror_figure_im.write(full_mirror_figure_file)
+
+    # The above file is the mirror_figure_image in the des template.
+    model = piff.Optical(template='des',atmo_type='Kolmogorov')
+    model_figure = model.mirror_figure_screen.table.f
+    np.testing.assert_array_equal(mirror_figure_im.array, model_figure)
+
+    # Usually mirror_figure_im is a file name, but you can also pass in an already read image.
+    model2 = piff.Optical(template='des', atmo_type='Kolmogorov',
+                          mirror_figure_im=mirror_figure_im)
+    model_figure = model2.mirror_figure_screen.table.f
+    np.testing.assert_array_equal(mirror_figure_im.array, model_figure)
+    prof1 = model.getProfile(zernike_coeff=[0,0,0,0,0], r0=0.15)
+    prof2 = model2.getProfile(zernike_coeff=[0,0,0,0,0], r0=0.15)
+    assert prof1 == prof2
 
+    # This is the old way that x and y were build for the UserScreen table.
+    # Make sure the new version is consistent.
+    mirror_figure_u = np.linspace(-mirror_figure_halfsize, mirror_figure_halfsize, num=512)
+    mirror_figure_v = np.linspace(-mirror_figure_halfsize, mirror_figure_halfsize, num=512)
+    np.testing.assert_allclose(model.mirror_figure_screen.table.x, mirror_figure_u, rtol=1.e-15)
+    np.testing.assert_allclose(model.mirror_figure_screen.table.y, mirror_figure_v, rtol=1.e-15)
+
+    # Overriding mirror_figure_scale will adjust the u,v values in the screen.
+    model3 = piff.Optical(template='des', atmo_type='Kolmogorov',
+                          mirror_figure_scale=ref_scale/2)
+    np.testing.assert_allclose(model3.mirror_figure_screen.table.x, mirror_figure_u/2, rtol=1.e-15)
+    np.testing.assert_allclose(model3.mirror_figure_screen.table.y, mirror_figure_v/2, rtol=1.e-15)
+
+    # Error if file not found.
+    with np.testing.assert_raises(ValueError):
+        model = piff.Optical(template='des', atmo_type='Kolmogorov',
+                             mirror_figure_im='invalid.fits')
 
 @timer
 def test_kolmogorov():
@@ -204,14 +278,6 @@ def test_kolmogorov():
     chi2 = np.std((star1.image - star2.image).array)
     assert chi2 != 0,'chi2 is zero!?'
 
-    # Usually mirror_figure_im is a file name, but you can also pass in an already read image.
-    mirror_figure_image = galsim.fits.read('input/DECam_236392_finegrid512_nm_uv.fits')
-    model2 = piff.Optical(template='des', atmo_type='Kolmogorov',
-                          mirror_figure_im=mirror_figure_image)
-    prof1 = model.getProfile(zernike_coeff=[0,0,0,0,0], r0=0.15)
-    prof2 = model2.getProfile(zernike_coeff=[0,0,0,0,0], r0=0.15)
-    assert prof1 == prof2
-
 
 @timer
 def test_vonkarman():
@@ -310,14 +376,15 @@ def test_disk():
     for key in model.opt_kwargs:
         assert key in model2.opt_kwargs, 'key %r missing from model2 opt_kwargs'%key
         if type(model.opt_kwargs[key])==np.ndarray:
-            np.testing.assert_almost_equal(model.opt_kwargs[key],model2.opt_kwargs[key],err_msg='key %r mismatch' % key)
+            np.testing.assert_almost_equal(model.opt_kwargs[key],
+                                           model2.opt_kwargs[key],
+                                           err_msg='key %r mismatch' % key)
         else:
             assert model.opt_kwargs[key] == model2.opt_kwargs[key], 'key %r mismatch'%key
     assert model.gsparams == model2.gsparams
-    for key in model.other_kwargs:
-        assert key in model2.other_kwargs, 'key %r missing from model2 other_kwargs'%key
-        assert model.other_kwargs[key] == model2.other_kwargs[key], 'key %r mismatch'%key
-    assert model.atmo_type == model2.atmo_type,'atmo_type mismatch'
+    assert model.sigma == model2.sigma
+    assert model.mirror_figure_screen == model2.mirror_figure_screen
+    assert model.atmo_type == model2.atmo_type
 
 #####
 # convenience functions
@@ -362,8 +429,9 @@ def make_empty_star(icen=500, jcen=700, ccdnum=28, params=None, stamp_size=40,
     test_init()
     test_optical()
     test_draw()
-    test_pupil_im(pupil_plane_file='input/DECam_pupil_512uv.fits')
-    test_pupil_im(pupil_plane_file='input/DECam_pupil_128uv.fits')
+    test_pupil_im(pupil_plane_file='DECam_pupil_512uv.fits')
+    test_pupil_im(pupil_plane_file='DECam_pupil_128uv.fits')
+    test_mirror_figure()
     test_kolmogorov()
     test_vonkarman()
     test_shearing()