Implement reconst_grey_2step_using_image()

Also added alpha parameter to runbsmem functions.
jsy1001 · Sep 12, 2018 · 2e81d63 · 2e81d63
1 parent 29c53af
commit 2e81d63
Showing 4 changed files with 119 additions and 25 deletions.
diff --git a/oirunner/priorimage.py b/oirunner/priorimage.py
@@ -16,6 +16,19 @@
 MAS_TO_DEG = 1/3600/1000
 
 
+def get_pixelsize(imagehdu: Union[fits.PrimaryHDU, fits.ImageHDU]) -> float:
+    """Return image pixel size in milliarcseconds."""
+    try:
+        cdelt1 = imagehdu.header['CDELT1']
+        cdelt2 = imagehdu.header['CDELT2']
+    except KeyError:
+        raise KeyError("CDELT1/2 keywords missing, pixelsize unknown")
+    if abs(cdelt1) != abs(cdelt2):
+        raise ValueError("Image pixels are not square " +
+                         "(CDELT1=%f, CDELT2=%f)" % (cdelt1, cdelt2))
+    return cdelt1 / MAS_TO_DEG
+
+
 def makesf(imagehdu: Union[fits.PrimaryHDU, fits.ImageHDU],
            fwhm: float, threshold: float) -> fits.PrimaryHDU:
     """Blur and threshold image for use as BSMEM prior model.
@@ -32,24 +45,16 @@ def makesf(imagehdu: Union[fits.PrimaryHDU, fits.ImageHDU],
       KeyError, ValueError
 
     """
-    # Read image
+    # Get image attributes
     dims = imagehdu.data.shape
-    try:
-        cdelt1 = imagehdu.header['CDELT1']
-        cdelt2 = imagehdu.header['CDELT2']
-    except KeyError:
-        raise KeyError("CDELT1/2 keywords missing, pixelsize unknown")
-    if abs(cdelt1) != abs(cdelt2):
-        raise ValueError("Image pixels are not square " +
-                         "(CDELT1=%f, CDELT2=%f)" % (cdelt1, cdelt2))
-    pixelsize = cdelt1 / MAS_TO_DEG
+    pixelsize = get_pixelsize(imagehdu)
     minvalue = imagehdu.data.min()
     maxvalue = imagehdu.data.max()
     logging.info('Image pixelsize = %f mas' % pixelsize)
     logging.info('Image min = %g' % minvalue)
     logging.info('Image max = %g' % maxvalue)
 
-    # Parameters:
+    # Initialize parameters
     sigma = fwhm / pixelsize / 2.3548
     lowest = threshold * maxvalue
     blank = 1e-8

diff --git a/oirunner/runbsmem.py b/oirunner/runbsmem.py
@@ -16,7 +16,7 @@
 
 from astropy.io import fits
 
-from .priorimage import makesf
+from .priorimage import get_pixelsize, makesf
 
 
 BSMEM = 'bsmem'
@@ -58,7 +58,8 @@ def run_bsmem(args: Sequence[str], fullstdout: str = None) -> None:
 def run_bsmem_using_model(datafile: str, outputfile: str, dim:
                           int, modeltype: int, modelwidth: float,
                           pixelsize: float = None,
-                          uvmax: float = None) -> None:
+                          uvmax: float = None,
+                          alpha: float = None) -> None:
     """Run bsmem using initial/prior model.
 
     Args:
@@ -69,6 +70,7 @@ def run_bsmem_using_model(datafile: str, outputfile: str, dim:
       modelwidth: Initial/prior image model width (mas).
       pixelsize:  Reconstructed image pixel size (mas).
       uvmax:      Maximum uv radius to select (waves).
+      alpha:      Regularization hyperparameter.
 
     """
     args = [BSMEM, '--noui',
@@ -81,13 +83,18 @@ def run_bsmem_using_model(datafile: str, outputfile: str, dim:
         args += ['--pixelsize=%f' % pixelsize]
     if uvmax is not None:
         args += ['--uvmax=%f' % uvmax]
+    if alpha is not None:
+        args += ['--autoalpha=3', '--alpha=%f' % alpha]
+    else:
+        args += ['--autoalpha=4']
     fullstdout = os.path.splitext(outputfile)[0] + '-out.txt'
     run_bsmem(args, fullstdout)
 
 
 def run_bsmem_using_image(datafile: str, outputfile: str, dim: int,
                           pixelsize: float, imagehdu: fits.PrimaryHDU,
-                          uvmax: float = None) -> None:
+                          uvmax: float = None,
+                          alpha: float = None) -> None:
     """Run bsmem using initial/prior image.
 
     Args:
@@ -97,6 +104,7 @@ def run_bsmem_using_image(datafile: str, outputfile: str, dim: int,
       pixelsize:  Reconstructed image pixel size (mas).
       imagehdu:   FITS HDU containing initial/prior image.
       uvmax:      Maximum uv radius to select (waves).
+      alpha:      Regularization hyperparameter.
 
     """
     tempimage = tempfile.NamedTemporaryFile(suffix='.fits', mode='wb',
@@ -111,6 +119,10 @@ def run_bsmem_using_image(datafile: str, outputfile: str, dim: int,
             '--sf=%s' % tempimage.name]
     if uvmax is not None:
         args += ['--uvmax=%f' % uvmax]
+    if alpha is not None:
+        args += ['--autoalpha=3', '--alpha=%f' % alpha]
+    else:
+        args += ['--autoalpha=4']
     fullstdout = os.path.splitext(outputfile)[0] + '-out.txt'
     run_bsmem(args, fullstdout)
     os.remove(tempimage.name)
@@ -121,7 +133,8 @@ def reconst_grey_basic(datafile: str,
                        dim: int = DEFAULT_DIM,
                        modeltype: int = DEFAULT_MT,
                        modelwidth: float = DEFAULT_MW,
-                       uvmax: float = None) -> None:
+                       uvmax: float = None,
+                       alpha: float = None) -> None:
     """Reconstruct a grey image by running bsmem once.
 
     Args:
@@ -131,11 +144,33 @@ def reconst_grey_basic(datafile: str,
       modeltype:  Initial/prior image model type (0-4).
       modelwidth: Initial/prior image model width (mas).
       uvmax:      Maximum uv radius to select (waves).
+      alpha:      Regularization hyperparameter.
 
     """
     run_bsmem_using_model(datafile, _get_outputfile(datafile, 1), dim,
                           modeltype, modelwidth,
-                          pixelsize=pixelsize, uvmax=uvmax)
+                          pixelsize=pixelsize, uvmax=uvmax, alpha=alpha)
+
+
+def reconst_grey_basic_using_image(datafile: str,
+                                   imagefile: str,
+                                   uvmax: float = None,
+                                   alpha: float = None) -> None:
+    """Reconstruct a grey image by running bsmem once using a prior image.
+
+    Args:
+      datafile:   Input OIFITS data filename.
+      imagefile:  Input initial/prior FITS image.
+      uvmax:      Maximum uv radius to select (waves).
+      alpha:      Regularization hyperparameter.
+
+    """
+    with fits.open(imagefile) as hdulist:
+        imagehdu = hdulist[0]
+        dim = imagehdu.data.shape[0]
+        pixelsize = get_pixelsize(imagehdu)
+        run_bsmem_using_image(datafile, _get_outputfile(datafile, 1), dim,
+                              pixelsize, imagehdu, uvmax=uvmax, alpha=alpha)
 
 
 def reconst_grey_2step(datafile: str,
@@ -144,6 +179,7 @@ def reconst_grey_2step(datafile: str,
                        modeltype: int = DEFAULT_MT,
                        modelwidth: float = DEFAULT_MW,
                        uvmax1: float = 1.1e8,
+                       alpha: float = None,
                        fwhm: float = 1.25,
                        threshold: float = 0.05) -> None:
     """Reconstruct a grey image by running bsmem twice.
@@ -155,15 +191,46 @@ def reconst_grey_2step(datafile: str,
       modeltype:  Initial/prior image model type for 1st run (0-4).
       modelwidth: Initial/prior image model width for 1st run (mas).
       uvmax1:     Maximum uv radius to select for 1st run (waves).
+      alpha:      Regularization hyperparameter for both runs.
       fwhm:       FWHM of Gaussian to convolve 1st run output with (mas).
       threshold:  Threshold (relative to peak) to apply to 1st run output.
 
     """
     # :TODO: intelligent defaults for uvmax1, fwhm?
     out1file = _get_outputfile(datafile, 1)
     run_bsmem_using_model(datafile, out1file, dim, modeltype, modelwidth,
-                          pixelsize=pixelsize, uvmax=uvmax1)
+                          pixelsize=pixelsize, uvmax=uvmax1, alpha=alpha)
     with fits.open(out1file) as hdulist:
         imagehdu = makesf(hdulist[0], fwhm, threshold)
     run_bsmem_using_image(datafile, _get_outputfile(datafile, 2),
-                          dim, pixelsize, imagehdu)
+                          dim, pixelsize, imagehdu, alpha=alpha)
+
+
+def reconst_grey_2step_using_image(datafile: str,
+                                   imagefile: str,
+                                   uvmax1: float = 1.1e8,
+                                   alpha: float = None,
+                                   fwhm: float = 1.25,
+                                   threshold: float = 0.05) -> None:
+    """Reconstruct a grey image by running bsmem twice using a prior image.
+
+    Args:
+      datafile:   Input OIFITS data filename.
+      imagefile:  Input initial/prior FITS image.
+      uvmax1:     Maximum uv radius to select for 1st run (waves).
+      alpha:      Regularization hyperparameter for both runs.
+      fwhm:       FWHM of Gaussian to convolve 1st run output with (mas).
+      threshold:  Threshold (relative to peak) to apply to 1st run output.
+
+    """
+    out1file = _get_outputfile(datafile, 1)
+    with fits.open(imagefile) as hdulist:
+        image1hdu = hdulist[0]
+        dim = image1hdu.data.shape[0]
+        pixelsize = get_pixelsize(image1hdu)
+        run_bsmem_using_image(datafile, out1file, dim,
+                              pixelsize, image1hdu, uvmax=uvmax1, alpha=alpha)
+    with fits.open(out1file) as hdulist:
+        image2hdu = makesf(hdulist[0], fwhm, threshold)
+    run_bsmem_using_image(datafile, _get_outputfile(datafile, 2),
+                          dim, pixelsize, image2hdu, alpha=alpha)
diff --git a/tests/gauss10.fits b/tests/gauss10.fits
diff --git a/tests/test_runbsmem.py b/tests/test_runbsmem.py
@@ -9,28 +9,50 @@
 except (OSError, CalledProcessError):
     HAVE_BSMEM = False
 
-from oirunner.runbsmem import reconst_grey_basic, reconst_grey_2step
+import oirunner.runbsmem as runbs
 
 DATAFILE = 'tests/2004contest1.oifits'
+IMAGEFILE = 'tests/gauss10.fits'
 
 
 class RunBsmemTestCase(unittest.TestCase):
 
     @unittest.skipUnless(HAVE_BSMEM, "requires bsmem")
     def test_grey_basic(self):
-        """Test basic grey reconstruction"""
+        """Test grey reconstruction"""
         with tempfile.TemporaryDirectory() as dirname:
             tempdatafile = os.path.join(dirname, os.path.basename(DATAFILE))
             copyfile(DATAFILE, tempdatafile)
-            reconst_grey_basic(tempdatafile)
-            reconst_grey_basic(tempdatafile, pixelsize=0.25)
-            reconst_grey_basic(tempdatafile, uvmax=1.1e8)
-            reconst_grey_basic(tempdatafile, pixelsize=0.25, uvmax=1.1e8)
+            runbs.reconst_grey_basic(tempdatafile)
+            runbs.reconst_grey_basic(tempdatafile, pixelsize=0.25)
+            runbs.reconst_grey_basic(tempdatafile, uvmax=1.1e8)
+            runbs.reconst_grey_basic(tempdatafile, alpha=4000.)
+            runbs.reconst_grey_basic(tempdatafile, pixelsize=0.25,
+                                     uvmax=1.1e8, alpha=4000.)
+
+    @unittest.skipUnless(HAVE_BSMEM, "requires bsmem")
+    def test_grey_basic_using_image(self):
+        """Test grey reconstruction using a prior image"""
+        with tempfile.TemporaryDirectory() as dirname:
+            tempdatafile = os.path.join(dirname, os.path.basename(DATAFILE))
+            copyfile(DATAFILE, tempdatafile)
+            runbs.reconst_grey_basic_using_image(tempdatafile, IMAGEFILE)
+            runbs.reconst_grey_basic(tempdatafile, uvmax=1.1e8)
+            runbs.reconst_grey_basic(tempdatafile, alpha=4000.)
+            runbs.reconst_grey_basic(tempdatafile, uvmax=1.1e8, alpha=4000.)
 
     @unittest.skipUnless(HAVE_BSMEM, "requires bsmem")
     def test_grey_2step(self):
         """Test two-step grey reconstruction"""
         with tempfile.TemporaryDirectory() as dirname:
             tempdatafile = os.path.join(dirname, os.path.basename(DATAFILE))
             copyfile(DATAFILE, tempdatafile)
-            reconst_grey_2step(tempdatafile, 0.25)
+            runbs.reconst_grey_2step(tempdatafile, 0.25)
+
+    @unittest.skipUnless(HAVE_BSMEM, "requires bsmem")
+    def test_grey_2step_using_image(self):
+        """Test two-step grey reconstruction using a prior image"""
+        with tempfile.TemporaryDirectory() as dirname:
+            tempdatafile = os.path.join(dirname, os.path.basename(DATAFILE))
+            copyfile(DATAFILE, tempdatafile)
+            runbs.reconst_grey_2step_using_image(tempdatafile, IMAGEFILE)