Fixes merging, improves f corona modeling

punchbowl/level1/flow.py

@@ -119,9 +119,9 @@
         # set up alignment mask
         observatory = int(data.meta["OBSCODE"].value)
         if observatory < 4:
-            alignment_mask = lambda x, y: (x > 100) * (x < 1900) * (y > 250) * (y < 1900)  # noqa: E731
+            alignment_mask = lambda x, y: (x > 100) * (x < 1900) * (y > 250) * (y < 1900)
         else:
-            alignment_mask = lambda x, y: (((x < 824) + (x > 1224)) * ((y < 824) + (y > 1224))  # noqa: E731
+            alignment_mask = lambda x, y: (((x < 824) + (x > 1224)) * ((y < 824) + (y > 1224))
                                            * (x > 100) * (x < 1900) * (y > 100) * (y < 1900))
         data = align_task(data, mask=alignment_mask)
 
@@ -197,9 +197,9 @@
         # set up alignment mask
         observatory = int(data.meta["OBSCODE"].value)
         if observatory < 4:
-            alignment_mask = lambda x, y: (x > 100) * (x < 1900) * (y > 250) * (y < 1900)  # noqa: E731
+            alignment_mask = lambda x, y: (x > 100) * (x < 1900) * (y > 250) * (y < 1900)
         else:
-            alignment_mask = lambda x, y: (((x < 824) + (x > 1224)) * ((y < 824) + (y > 1224))  # noqa: E731
+            alignment_mask = lambda x, y: (((x < 824) + (x > 1224)) * ((y < 824) + (y > 1224))
                                            * (x > 100) * (x < 1900) * (y > 100) * (y < 1900))
 
         data = align_task(data, mask=alignment_mask)

punchbowl/level2/flow.py

@@ -1,3 +1,4 @@
+from glob import glob
 from datetime import datetime
 
 import numpy as np
@@ -131,3 +132,8 @@
 
     logger.info("ending level quickPUNCH core flow")
     return [output_data_mosaic, output_data_nfi]
+
+
+if __name__ == "__main__":
+    filenames = glob("/Users/jhughes/new_results/nov25-1026/cr/*.fits")
+    out = levelq_core_flow(filenames)

punchbowl/level2/merge.py

@@ -16,11 +16,12 @@
         selected_images = [d for d in data if d is not None and d.meta["POLAR"].value == polarization]
         if len(selected_images) > 0:
             reprojected_data = np.stack([d.data for d in selected_images], axis=-1)
-            reprojected_weights = np.stack([1/np.square(d.uncertainty.array) for d in selected_images], axis=-1)
+            reprojected_uncertainties = np.stack([d.uncertainty.array for d in selected_images], axis=-1)
+            reprojected_uncertainties[reprojected_uncertainties <= 0] = np.inf
+            reprojected_uncertainties[np.isinf(reprojected_uncertainties)] = 1E64
+            reprojected_uncertainties[reprojected_data <= 0] = np.inf
 
-            reprojected_weights[reprojected_weights <= 0] = 1E-16
-            reprojected_weights[np.isinf(reprojected_weights)] = 1E16
-            reprojected_weights[np.isnan(reprojected_weights)] = 1E-16
+            reprojected_weights = 1 / np.square(reprojected_uncertainties)
 
             trefoil_data_layers.append(np.nansum(reprojected_data * reprojected_weights, axis=2) /
                                        np.nansum(reprojected_weights, axis=2))
@@ -40,19 +41,20 @@
 def merge_many_clear_task(data: list[NDCube | None], trefoil_wcs: WCS) -> NDCube:
     """Merge many task and carefully combine uncertainties."""
     trefoil_data_layers, trefoil_uncertainty_layers = [], []
-    selected_images = [d.data for d in data if d is not None]
+    selected_images = [d for d in data if d is not None]
 
     if len(selected_images) > 0:
-        reprojected_data = np.stack(selected_images, axis=-1)
-        reprojected_weights = np.stack([1/np.square(d.uncertainty.array) for d in data], axis=-1)
+        reprojected_data = np.stack([d.data for d in selected_images], axis=-1)
+        reprojected_uncertainties = np.stack([d.uncertainty.array for d in selected_images], axis=-1)
+        reprojected_uncertainties[reprojected_uncertainties <= 0] = np.inf
+        reprojected_uncertainties[np.isinf(reprojected_uncertainties)] = 1E64
+        reprojected_uncertainties[reprojected_data <= 0] = np.inf
 
-        reprojected_weights[reprojected_weights <= 0] = 1E-16
-        reprojected_weights[np.isinf(reprojected_weights)] = 1E16
-        reprojected_weights[np.isnan(reprojected_weights)] = 1E-16
+        reprojected_weights = 1/np.square(reprojected_uncertainties)
 
-        trefoil_data_layers.append(np.nansum(reprojected_data * reprojected_weights, axis=2) /
-                                   np.nansum(reprojected_weights, axis=2))
-        trefoil_uncertainty_layers.append(1/np.nansum(np.sqrt(reprojected_weights), axis=2))
+        trefoil_data_layers.append(np.nansum(reprojected_data * reprojected_weights, axis=-1) /
+                                   np.nansum(reprojected_weights, axis=-1))
+        trefoil_uncertainty_layers.append(1/np.sqrt(np.nansum(reprojected_weights, axis=-1)))
     else:
         trefoil_data_layers.append(np.zeros((4096, 4096)))
         trefoil_uncertainty_layers.append(np.zeros((4096, 4096))-999)
@@ -61,8 +63,8 @@
     output_meta["DATE-OBS"] = data[0].meta["DATE-OBS"].value  # TODO: do this better and fill rest of meta
 
     return NDCube(
-        data=np.stack(trefoil_data_layers, axis=0),
-        uncertainty=StdDevUncertainty(np.stack(trefoil_uncertainty_layers, axis=0)),
+        data=np.stack(trefoil_data_layers, axis=0).squeeze(),
+        uncertainty=StdDevUncertainty(np.stack(trefoil_uncertainty_layers, axis=0).squeeze()),
         wcs=trefoil_wcs,
         meta=output_meta,
     )

punchbowl/level2/tests/test_merge.py

@@ -34,13 +34,10 @@ def test_merge_many_task(sample_data_list):
     trefoil_wcs = WCS("level2/data/trefoil_hdr.fits")
     trefoil_wcs.wcs.ctype = "HPLN-ARC", "HPLT-ARC"
 
-    pd_list = sample_data_list
-
-
-    output_punchdata = merge_many_polarized_task.fn(pd_list, trefoil_wcs)
+    output_punchdata = merge_many_polarized_task.fn(sample_data_list, trefoil_wcs)
     assert isinstance(output_punchdata, NDCube)
     assert output_punchdata.data.shape == (3, 4096, 4096)
     assert np.allclose(output_punchdata.data[0, 100:300, 300:400], 1)
-    assert np.allclose(output_punchdata.data[0, 500:700, 800:900], 500)
+    assert np.allclose(output_punchdata.data[0, 500:700, 800:900], 1)
     assert np.allclose(output_punchdata.uncertainty.array[0, 100:300, 300:400], 1)
-    assert np.allclose(output_punchdata.uncertainty.array[0, 500:700, 800:900], 0)
+    assert np.allclose(output_punchdata.uncertainty.array[0, 500:700, 800:900], 1)

punchbowl/level3/f_corona_model.py

@@ -1,14 +1,15 @@
 from datetime import datetime
 
 import numpy as np
+from dateutil.parser import parse as parse_datetime_str
 from ndcube import NDCube
 from numpy.polynomial import polynomial
 from prefect import flow, get_run_logger
 from quadprog import solve_qp
 from scipy.interpolate import griddata
 
 from punchbowl.data import NormalizedMetadata, load_ndcube_from_fits
-from punchbowl.data.wcs import load_trefoil_wcs
+from punchbowl.data.wcs import load_quickpunch_mosaic_wcs, load_trefoil_wcs
 from punchbowl.exceptions import InvalidDataError
 from punchbowl.prefect import punch_task
 
@@ -35,20 +36,19 @@
 
     """
     c = np.transpose(input_vals)
+    cube_is_good = np.isfinite(cube)
+    num_inputs = np.sum(cube_is_good, axis=0)
 
     solution = np.zeros((input_vals.shape[1], cube.shape[1], cube.shape[2]))
-    num_inputs = np.zeros((cube.shape[1], cube.shape[2]))
     for i in range(cube.shape[1]):
         for j in range(cube.shape[2]):
-            time_series = cube[:, i, j]
-            is_good = np.isfinite(time_series)
-            time_series = time_series[is_good]
-            c_iter = c[:, is_good]
-            g_iter = np.matmul(c_iter, c_iter.T)
-            num_inputs[i, j] = np.sum(is_good)
+            is_good = cube_is_good[:, i, j]
+            time_series = cube[:, i, j][is_good]
             if time_series.size < n_nonnan_required:
                 this_solution = np.zeros(input_vals.shape[1])
             else:
+                c_iter = c[:, is_good]
+                g_iter = np.matmul(c_iter, c_iter.T)
                 a = np.matmul(c_iter, time_series)
                 try:
                     this_solution = solve_qp(g_iter, a, c_iter, time_series)[0]
@@ -59,8 +59,9 @@
     return np.asarray(solution), num_inputs
 
 def model_fcorona_for_cube(xt: np.ndarray,
+                           reference_xt: float,
                           cube: np.ndarray,
-                          min_brightness: float = 1E-16,
+                          min_brightness: float = 1E-18,
                           smooth_level: float | None = 1,
                           return_full_curves: bool=False,
                           ) -> tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -71,6 +72,8 @@
     ----------
     xt : np.ndarray
         time array
+    reference_xt: float
+        timestamp to evaluate the model for
     cube : np.ndarray
         observation array
     min_brightness: float
@@ -96,23 +99,24 @@
     """
     cube[cube < min_brightness] = np.nan
     if smooth_level is not None:
-        average = np.nanmean(cube, axis=0)
-        std = np.nanstd(cube, axis=0)
-        a, b, c = np.where(cube[:, ...] > (average + (smooth_level * std)))
-        cube[a, b, c] = average[b, c]
+        center = np.nanmedian(cube, axis=0)
+        width = np.diff(np.nanpercentile(cube, [25, 75], axis=0), axis=0).squeeze()
+        a, b, c = np.where(cube[:, ...] > (center + (smooth_level * width)))
+        cube[a, b, c] = center[b, c]
 
-        a, b, c = np.where(cube[:, ...] < (average - (smooth_level * std)))
-        cube[a, b, c] = average[b, c]
+        a, b, c = np.where(cube[:, ...] < (center - (smooth_level * width)))
+        cube[a, b, c] = center[b, c]
 
     xt = np.array(xt)
+    reference_xt -= xt[0]
     xt -= xt[0]
 
     input_array = np.c_[np.power(xt, 3), np.square(xt), xt, np.ones(len(xt))]
     coefficients, counts = solve_qp_cube(input_array, -cube)
     coefficients *= -1
     if return_full_curves:
         return polynomial.polyval(xt, coefficients[::-1, :, :]).transpose((2, 0, 1)), counts, cube
-    return polynomial.polyval(xt[len(xt)//2], coefficients[::-1, :, :]), counts
+    return polynomial.polyval(reference_xt, coefficients[::-1, :, :]), counts
 
 
 def fill_nans_with_interpolation(image: np.ndarray) -> np.ndarray:
@@ -125,10 +129,16 @@
     return griddata((x, y), known_values, (grid_x, grid_y), method="cubic")
 
 @flow(log_prints=True)
-def construct_polarized_f_corona_model(filenames: list[str], smooth_level: float = 3.0) -> list[NDCube]:
+def construct_polarized_f_corona_model(filenames: list[str], smooth_level: float = 3.0,
+                                       reference_time: str | None = None) -> list[NDCube]:
     """Construct a full F corona model."""
     logger = get_run_logger()
 
+    if reference_time is None:
+        reference_time = datetime.now()
+    elif isinstance(reference_time, str):
+        reference_time = parse_datetime_str(reference_time)
+
     trefoil_wcs, trefoil_shape = load_trefoil_wcs()
 
     logger.info("construct_f_corona_background started")
@@ -142,8 +152,8 @@
     data_shape = (3, *trefoil_shape)
 
     number_of_data_frames = len(filenames)
-    data_cube = np.empty((number_of_data_frames, data_shape), dtype=float)
-    uncertainty_cube = np.empty((number_of_data_frames, data_shape), dtype=float)
+    data_cube = np.empty((number_of_data_frames, *data_shape), dtype=float)
+    uncertainty_cube = np.empty((number_of_data_frames, *data_shape), dtype=float)
 
     meta_list = []
     obs_times = []
@@ -157,21 +167,28 @@
         meta_list.append(data_object.meta)
     logger.info("ending data loading")
 
-    m_model_fcorona, _ = model_fcorona_for_cube(obs_times, data_cube[0], smooth_level=smooth_level)
-    m_model_fcorona.data[m_model_fcorona.data==0] = np.nan
-    m_model_fcorona = fill_nans_with_interpolation(m_model_fcorona.data)
-
-    z_model_fcorona, _ = model_fcorona_for_cube(obs_times, data_cube[1], smooth_level=smooth_level)
-    z_model_fcorona.data[z_model_fcorona.data==0] = np.nan
-    z_model_fcorona = fill_nans_with_interpolation(z_model_fcorona.data)
-
-    p_model_fcorona, _ = model_fcorona_for_cube(obs_times, data_cube[2], smooth_level=smooth_level)
-    p_model_fcorona.data[p_model_fcorona.data==0] = np.nan
-    p_model_fcorona = fill_nans_with_interpolation(p_model_fcorona.data)
+    reference_xt = reference_time.timestamp()
+    m_model_fcorona, _ = model_fcorona_for_cube(obs_times, reference_xt,
+                                                data_cube[:, 0, :, :], smooth_level=smooth_level)
+    m_model_fcorona[m_model_fcorona==0] = np.nan
+    m_model_fcorona = fill_nans_with_interpolation(m_model_fcorona)
+
+    z_model_fcorona, _ = model_fcorona_for_cube(obs_times,
+                                                reference_xt,
+                                                data_cube[:, 1, :, :],
+                                                smooth_level=smooth_level)
+    z_model_fcorona[z_model_fcorona==0] = np.nan
+    z_model_fcorona = fill_nans_with_interpolation(z_model_fcorona)
+
+    p_model_fcorona, _ = model_fcorona_for_cube(obs_times,
+                                                reference_xt,
+                                                data_cube[:, 2, :, :],
+                                                smooth_level=smooth_level)
+    p_model_fcorona[p_model_fcorona==0] = np.nan
+    p_model_fcorona = fill_nans_with_interpolation(p_model_fcorona)
 
     meta = NormalizedMetadata.load_template("PFM", "3")
-    meta["DATE-OBS"] =  str(datetime(2024, 12, 1, 12, 0, 0,
-                                     tzinfo=datetime.timezone.utc))
+    meta["DATE-OBS"] = str(reference_time)
     output_cube = NDCube(data=np.stack([m_model_fcorona,
                                                z_model_fcorona,
                                                p_model_fcorona], axis=0),
@@ -278,3 +295,56 @@
 def create_empty_f_background_model(data_object: NDCube) -> np.ndarray:
     """Create an empty background model."""
     return np.zeros_like(data_object.data)
+
+
+@flow(log_prints=True)
+def construct_qp_f_corona_model(filenames: list[str], smooth_level: float = 3.0,
+                                       reference_time: str | None = None) -> list[NDCube]:
+    """Construct QuickPUNCH F corona model."""
+    logger = get_run_logger()
+
+    if reference_time is None:
+        reference_time = datetime.now()
+    elif isinstance(reference_time, str):
+        reference_time = parse_datetime_str(reference_time)
+
+    trefoil_wcs, trefoil_shape = load_quickpunch_mosaic_wcs()
+
+    logger.info("construct_f_corona_background started")
+
+    if len(filenames) == 0:
+        msg = "Require at least one input file"
+        raise ValueError(msg)
+
+    filenames.sort()
+
+    data_shape = trefoil_shape
+
+    number_of_data_frames = len(filenames)
+    data_cube = np.empty((number_of_data_frames, *data_shape), dtype=float)
+    uncertainty_cube = np.empty((number_of_data_frames, *data_shape), dtype=float)
+
+    meta_list = []
+    obs_times = []
+
+    logger.info("beginning data loading")
+    for i, address_out in enumerate(filenames):
+        data_object = load_ndcube_from_fits(address_out)
+        data_cube[i, ...] = data_object.data
+        uncertainty_cube[i, ...] = data_object.uncertainty.array
+        obs_times.append(data_object.meta.datetime.timestamp())
+        meta_list.append(data_object.meta)
+    logger.info("ending data loading")
+
+    reference_xt = reference_time.timestamp()
+    model_fcorona, _ = model_fcorona_for_cube(obs_times, reference_xt, data_cube, smooth_level=smooth_level)
+    model_fcorona[model_fcorona<=0] = np.nan
+    model_fcorona = fill_nans_with_interpolation(model_fcorona)
+
+    meta = NormalizedMetadata.load_template("CFM", "Q")
+    meta["DATE-OBS"] = str(reference_time)
+    output_cube = NDCube(data=model_fcorona.squeeze(),
+                                meta=meta,
+                                wcs=trefoil_wcs)
+
+    return [output_cube]

punchbowl/level3/stellar.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import remove_starfield
+from dateutil.parser import parse as parse_datetime_str
 from ndcube import NDCube
 from prefect import flow, get_run_logger
 from remove_starfield import ImageHolder, ImageProcessor, Starfield
@@ -30,9 +31,16 @@
         filenames: list[str],
         n_sigma: float = 5,
         map_scale: float = 0.01,
-        target_mem_usage: float = 1000) -> [NDCube, NDCube]:
+        target_mem_usage: float = 1000,
+        reference_time: datetime | None = None) -> [NDCube, NDCube]:
     """Create a background starfield_bg map from a series of PUNCH images over a long period of time."""
     logger = get_run_logger()
+
+    if reference_time is None:
+        reference_time = datetime.now()
+    elif isinstance(reference_time, str):
+        reference_time = parse_datetime_str(reference_time)
+
     logger.info("construct_starfield_background started")
 
     # create an empty array to fill with data
@@ -85,28 +93,18 @@
         target_mem_usage=target_mem_usage)
     logger.info("Ending p starfield")
 
-
     # create an output PUNCHdata object
     logger.info("Preparing to create outputs")
 
     meta = NormalizedMetadata.load_template("PSM", "3")
-    meta["DATE-OBS"] = str(datetime(2024, 8, 1, 12, 0, 0,
-                                    tzinfo=datetime.timezone.utc))
-    out_wcs, _ = calculate_helio_wcs_from_celestial(starfield_m.wcs, meta.astropy_time, starfield_m.starfield.shape)
-    output_before = NDCube(np.stack([starfield_m.starfield, starfield_z.starfield, starfield_p.starfield], axis=0),
-                    wcs=out_wcs, meta=meta)
-    output_before.meta.history.add_now("LEVEL3-starfield_background", "constructed starfield_bg model")
-
-    meta = NormalizedMetadata.load_template("PSM", "3")
-    meta["DATE-OBS"] = str(datetime(2024, 12, 1, 12, 0, 0,
-                                    tzinfo=datetime.timezone.utc))
+    meta["DATE-OBS"] = reference_time
     out_wcs, _ = calculate_helio_wcs_from_celestial(starfield_m.wcs, meta.astropy_time, starfield_m.starfield.shape)
-    output_after = NDCube(np.stack([starfield_m.starfield, starfield_z.starfield, starfield_p.starfield], axis=0),
+    output = NDCube(np.stack([starfield_m.starfield, starfield_z.starfield, starfield_p.starfield], axis=0),
                     wcs=out_wcs, meta=meta)
-    output_after.meta.history.add_now("LEVEL3-starfield_background", "constructed starfield_bg model")
+    output.meta.history.add_now("LEVEL3-starfield_background", "constructed starfield_bg model")
 
     logger.info("construct_starfield_background finished")
-    return [output_before, output_after]
+    return [output]
 
 
 @punch_task