CHG: using header when writing fits files with pyralysis

miguelcarcamov · Oct 30, 2022 · 2e14af9 · 2e14af9
1 parent 1f8a5ee
commit 2e14af9
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/snow/imaging/gpuvmem.py b/src/snow/imaging/gpuvmem.py
@@ -214,6 +214,7 @@ def __restore_pyra(
         )
         with fits.open(model_fits) as hdu_model:
             im_model = hdu_model[0].data
+            header_model = hdu_model[0].header
         pix_scale = Quantity(self.cell)
 
         imsize = [self.M, self.N]
@@ -270,7 +271,7 @@ def __restore_pyra(
         )
         image = dirty_mapper.transform()[0].data[0].compute()
 
-        fits_io.write(image, output_name=residual_casa_image + ".fits")
+        fits_io.write(image, output_name=residual_casa_image + ".fits", header=header_model)
 
         psf = dataset.psf[0]
 
@@ -286,7 +287,7 @@ def __restore_pyra(
 
         im_restored = im_model_convolved + image
 
-        fits_io.write(im_restored, output_name=restored_image + ".fits")
+        fits_io.write(im_restored, output_name=restored_image + ".fits", header=header_model)
 
         return residual_casa_image + ".fits", restored_image + ".fits"