Skip to content

Commit

Permalink
Updates for V4 RFR2 (#346)
Browse files Browse the repository at this point in the history 
* use isoformat for date-obs

* fix timestamps

* fix the datetimes in l2 🐸

* check that dates are set

* toggle alignment

* fix lq assignment of date

* fix dates in f corona

* fix wcs shape

* improve starfield modeling

* allow missing data in merge

* remove silly import from build

* pins matplotlib version

* limit starfield number of procs

* change despike params

* manually set timeout_seconds to see if it helps

* fix starfield wcs

---------

Co-authored-by: Chris Lowder <[email protected]>
  • Loading branch information
@jmbhughes @lowderchris
jmbhughes and lowderchris authored Dec 19, 2024
1 parent e2ab02a commit 3b3b496
Show file tree
Hide file tree
Showing 8 changed files with 194 additions and 76 deletions.
20 changes: 10 additions & 10 deletions punchbowl/level1/flow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,9 +51,9 @@ def level1_core_flow(
read_noise_level: float = 17,
bitrate_signal: int = 16,
quartic_coefficient_path: str | pathlib.Path | None = None,
despike_unsharp_size: int = 2,
despike_unsharp_size: int = 1,
despike_method: str = "median",
despike_alpha: float = 50,
despike_alpha: float = 10,
despike_dilation: int = 0,
exposure_time: float = 49 * 1000,
readout_line_time: float = 163/2148,
Expand All @@ -64,7 +64,7 @@ def level1_core_flow(
deficient_pixel_required_good_count: int = 3,
deficient_pixel_max_window_size: int = 10,
psf_model_path: str | None = None,
alignment_mask: Callable | None = None,
alignment_mask: Callable | None = None, # noqa: ARG001
output_filename: list[str] | None = None,
) -> list[NDCube]:
"""Core flow for level 1."""
Expand Down Expand Up @@ -117,13 +117,13 @@ def level1_core_flow(
data = correct_psf_task(data, psf_model_path)

# 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)
else:
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)
# observatory = int(data.meta["OBSCODE"].value) # noqa: ERA001
# if observatory < 4:
# alignment_mask = lambda x, y: (x > 100) * (x < 1900) * (y > 250) * (y < 1900) # noqa: ERA001
# else: # noqa: ERA001
# 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) # noqa: ERA001

# Repackage data with proper metadata
product_code = data.meta["TYPECODE"].value + data.meta["OBSCODE"].value
Expand Down
117 changes: 86 additions & 31 deletions punchbowl/level2/flow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,9 @@ def level2_core_flow(data_list: list[str] | list[NDCube],

if data_list:
# order the data list so it can be processed properly
output_dateobs = str(average_datetime([d.meta.datetime for d in data_list]))
output_datebeg = str(min([d.meta.datetime for d in data_list]))
output_dateend = str(max([d.meta.datetime for d in data_list]))
output_dateobs = average_datetime([d.meta.datetime for d in data_list]).isoformat()
output_datebeg = min([d.meta.datetime for d in data_list]).isoformat()
output_dateend = max([d.meta.datetime for d in data_list]).isoformat()

ordered_data_list: list[NDCube | None] = [None for _ in range(len(ORDER))]
ordered_voters: list[list[str]] = [[] for _ in range(len(ORDER))]
Expand All @@ -64,8 +64,8 @@ def level2_core_flow(data_list: list[str] | list[NDCube],
for cube, this_voter_filenames in zip(data_list, ordered_voters, strict=True)]
output_data = merge_many_polarized_task(data_list, trefoil_wcs)
else:
output_dateobs = str(datetime.now())
output_datebeg= output_dateobs
output_dateobs = datetime.now().isoformat()
output_datebeg = output_dateobs
output_dateend = output_datebeg

output_data = NDCube(
Expand All @@ -75,6 +75,8 @@ def level2_core_flow(data_list: list[str] | list[NDCube],
meta=NormalizedMetadata.load_template("PTM", "2"),
)

output_data.meta["DATE"] = datetime.now().isoformat()
output_data.meta["DATE-AVG"] = output_dateobs
output_data.meta["DATE-OBS"] = output_dateobs
output_data.meta["DATE-BEG"] = output_datebeg
output_data.meta["DATE-END"] = output_dateend
Expand All @@ -97,33 +99,86 @@ def levelq_core_flow(data_list: list[str] | list[NDCube],

data_list = [load_image_task(d) if isinstance(d, str) else d for d in data_list]

ordered_data_list: list[NDCube | None] = [None for _ in range(len(ORDER_QP))]
for i, order_element in enumerate(ORDER_QP):
for data_element in data_list:
typecode = data_element.meta["TYPECODE"].value
obscode = data_element.meta["OBSCODE"].value
if typecode == order_element[:2] and obscode == order_element[2]:
ordered_data_list[i] = data_element
logger.info("Ordered files are "
f"{[get_base_file_name(cube) if cube is not None else None for cube in ordered_data_list]}")

quickpunch_mosaic_wcs, quickpunch_mosaic_shape = load_quickpunch_mosaic_wcs()
quickpunch_nfi_wcs, quickpunch_nfi_shape = load_quickpunch_nfi_wcs()

data_list_mosaic = reproject_many_flow(ordered_data_list, quickpunch_mosaic_wcs, quickpunch_mosaic_shape)
data_list_nfi = reproject_many_flow(ordered_data_list[-1:], quickpunch_nfi_wcs, quickpunch_nfi_shape)
output_data_mosaic = merge_many_clear_task(data_list_mosaic, quickpunch_mosaic_wcs)
output_data_mosaic = set_spacecraft_location_to_earth(output_data_mosaic)

output_meta_nfi = NormalizedMetadata.load_template("CNN", "Q")
output_meta_nfi["DATE-OBS"] = data_list_nfi[0].meta["DATE-OBS"].value
output_data_nfi = NDCube(
data=data_list_nfi[0].data,
uncertainty=StdDevUncertainty(data_list_nfi[0].uncertainty.array),
wcs=quickpunch_nfi_wcs,
meta=output_meta_nfi,
if data_list:
ordered_data_list: list[NDCube | None] = [None for _ in range(len(ORDER_QP))]
for i, order_element in enumerate(ORDER_QP):
for data_element in data_list:
typecode = data_element.meta["TYPECODE"].value
obscode = data_element.meta["OBSCODE"].value
if typecode == order_element[:2] and obscode == order_element[2]:
ordered_data_list[i] = data_element
logger.info("Ordered files are "
f"{[get_base_file_name(cube) if cube is not None else None for cube in ordered_data_list]}")

quickpunch_mosaic_wcs, quickpunch_mosaic_shape = load_quickpunch_mosaic_wcs()
quickpunch_nfi_wcs, quickpunch_nfi_shape = load_quickpunch_nfi_wcs()

data_list_mosaic = reproject_many_flow(ordered_data_list, quickpunch_mosaic_wcs, quickpunch_mosaic_shape)
output_dateobs = average_datetime([d.meta.datetime for d in data_list_mosaic]).isoformat()
output_datebeg = min([d.meta.datetime for d in data_list_mosaic]).isoformat()
output_dateend = max([d.meta.datetime for d in data_list_mosaic]).isoformat()

data_list_nfi = reproject_many_flow(ordered_data_list[-1:], quickpunch_nfi_wcs, quickpunch_nfi_shape)
output_dateobs_nfi = average_datetime([d.meta.datetime for d in data_list_nfi]).isoformat()
output_datebeg_nfi = min([d.meta.datetime for d in data_list_nfi]).isoformat()
output_dateend_nfi = max([d.meta.datetime for d in data_list_nfi]).isoformat()

output_data_mosaic = merge_many_clear_task(data_list_mosaic, quickpunch_mosaic_wcs)

output_data_mosaic.meta["DATE"] = datetime.now().isoformat()
output_data_mosaic.meta["DATE-AVG"] = output_dateobs
output_data_mosaic.meta["DATE-OBS"] = output_dateobs
output_data_mosaic.meta["DATE-BEG"] = output_datebeg
output_data_mosaic.meta["DATE-END"] = output_dateend
output_data_mosaic = set_spacecraft_location_to_earth(output_data_mosaic)

output_meta_nfi = NormalizedMetadata.load_template("CNN", "Q")
output_meta_nfi["DATE-OBS"] = data_list_nfi[0].meta["DATE-OBS"].value
output_data_nfi = NDCube(
data=data_list_nfi[0].data,
uncertainty=StdDevUncertainty(data_list_nfi[0].uncertainty.array),
wcs=quickpunch_nfi_wcs,
meta=output_meta_nfi,
)
output_data_nfi.meta["DATE"] = datetime.now().isoformat()
output_data_nfi.meta["DATE-AVG"] = output_dateobs_nfi
output_data_nfi.meta["DATE-OBS"] = output_dateobs_nfi
output_data_nfi.meta["DATE-BEG"] = output_datebeg_nfi
output_data_nfi.meta["DATE-END"] = output_dateend_nfi
output_data_nfi = set_spacecraft_location_to_earth(output_data_nfi)
else:
output_dateobs = datetime.now().isoformat()
output_datebeg = output_dateobs
output_dateend = output_datebeg

quickpunch_mosaic_wcs, quickpunch_mosaic_shape = load_quickpunch_mosaic_wcs()
quickpunch_nfi_wcs, quickpunch_nfi_shape = load_quickpunch_nfi_wcs()
output_data_mosaic = NDCube(
data=np.zeros(quickpunch_mosaic_shape),
uncertainty=StdDevUncertainty(np.zeros(quickpunch_mosaic_shape)),
wcs=quickpunch_mosaic_wcs,
meta=NormalizedMetadata.load_template("CTM", "Q"),
)
output_data_mosaic.meta["DATE"] = datetime.now().isoformat()
output_data_mosaic.meta["DATE-AVG"] = output_dateobs
output_data_mosaic.meta["DATE-OBS"] = output_dateobs
output_data_mosaic.meta["DATE-BEG"] = output_datebeg
output_data_mosaic.meta["DATE-END"] = output_dateend
output_data_mosaic = set_spacecraft_location_to_earth(output_data_mosaic)

output_data_nfi = NDCube(
data=np.zeros(quickpunch_nfi_shape),
uncertainty=StdDevUncertainty(np.zeros(quickpunch_nfi_shape)),
wcs=quickpunch_nfi_wcs,
meta=NormalizedMetadata.load_template("CNN", "Q"),
)
output_data_nfi = set_spacecraft_location_to_earth(output_data_nfi)

output_data_nfi.meta["DATE"] = datetime.now().isoformat()
output_data_nfi.meta["DATE-AVG"] = output_dateobs
output_data_nfi.meta["DATE-OBS"] = output_dateobs
output_data_nfi.meta["DATE-BEG"] = output_dateobs
output_data_nfi.meta["DATE-END"] = output_dateobs
output_data_nfi = set_spacecraft_location_to_earth(output_data_nfi)


if output_filename is not None:
Expand Down
16 changes: 12 additions & 4 deletions punchbowl/level2/merge.py
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@

import astropy
import numpy as np
from astropy.nddata import StdDevUncertainty
from astropy.wcs import WCS
from ndcube import NDCube

from punchbowl.data import NormalizedMetadata
from punchbowl.prefect import punch_task
from punchbowl.util import average_datetime


@punch_task
Expand All @@ -30,11 +31,19 @@ def merge_many_polarized_task(data: list[NDCube | None], trefoil_wcs: WCS) -> ND
trefoil_data_layers.append(np.zeros((4096, 4096)))
trefoil_uncertainty_layers.append(np.zeros((4096, 4096))-999)

output_meta = NormalizedMetadata.load_template("PTM", "2")
output_dateobs = average_datetime([d.meta.datetime for d in data if d is not None]).isoformat()
output_datebeg = min([d.meta.datetime for d in data if d is not None]).isoformat()
output_dateend = max([d.meta.datetime for d in data if d is not None]).isoformat()
output_meta["DATE-OBS"] = output_dateobs
output_meta["DATE-BEG"] = output_datebeg
output_meta["DATE-END"] = output_dateend
trefoil_3d_wcs = astropy.wcs.utils.add_stokes_axis_to_wcs(trefoil_wcs, 2)
return NDCube(
data=np.stack(trefoil_data_layers, axis=0),
uncertainty=StdDevUncertainty(np.stack(trefoil_uncertainty_layers, axis=0)),
wcs=trefoil_wcs,
meta=NormalizedMetadata.load_template("PTM", "2"),
wcs=trefoil_3d_wcs,
meta=output_meta,
)

@punch_task
Expand All @@ -60,7 +69,6 @@ def merge_many_clear_task(data: list[NDCube | None], trefoil_wcs: WCS) -> NDCube
trefoil_uncertainty_layers.append(np.zeros((4096, 4096))-999)

output_meta = NormalizedMetadata.load_template("CTM", "Q")
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).squeeze(),
Expand Down
18 changes: 16 additions & 2 deletions punchbowl/level3/f_corona_model.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from datetime import datetime

import astropy
import numpy as np
from dateutil.parser import parse as parse_datetime_str
from ndcube import NDCube
Expand All @@ -9,6 +10,7 @@
from scipy.interpolate import griddata

from punchbowl.data import NormalizedMetadata, load_ndcube_from_fits
from punchbowl.data.meta import set_spacecraft_location_to_earth
from punchbowl.data.wcs import load_quickpunch_mosaic_wcs, load_trefoil_wcs
from punchbowl.exceptions import InvalidDataError
from punchbowl.prefect import punch_task
Expand Down Expand Up @@ -160,13 +162,17 @@ def construct_polarized_f_corona_model(filenames: list[str], smooth_level: float
obs_times = []

logger.info("beginning data loading")
dates = []
for i, address_out in enumerate(filenames):
data_object = load_ndcube_from_fits(address_out)
dates.append(data_object.meta.datetime)
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")
output_datebeg = min(dates).isoformat()
output_dateend = max(dates).isoformat()

reference_xt = reference_time.timestamp()
m_model_fcorona, _ = model_fcorona_for_cube(obs_times, reference_xt,
Expand All @@ -189,12 +195,20 @@ def construct_polarized_f_corona_model(filenames: list[str], smooth_level: float
p_model_fcorona = fill_nans_with_interpolation(p_model_fcorona)

meta = NormalizedMetadata.load_template("PFM", "3")
meta["DATE-OBS"] = str(reference_time)
meta["DATE"] = datetime.now().isoformat()
meta["DATE-AVG"] = reference_time.isoformat()
meta["DATE-OBS"] = reference_time.isoformat()

meta["DATE-BEG"] = output_datebeg
meta["DATE-END"] = output_dateend
trefoil_3d_wcs = astropy.wcs.utils.add_stokes_axis_to_wcs(trefoil_wcs, 2)

output_cube = NDCube(data=np.stack([m_model_fcorona,
z_model_fcorona,
p_model_fcorona], axis=0),
meta=meta,
wcs=trefoil_wcs)
wcs=trefoil_3d_wcs)
output_cube = set_spacecraft_location_to_earth(output_cube)

return [output_cube]

Expand Down
14 changes: 12 additions & 2 deletions punchbowl/level3/flow.py
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
from datetime import datetime

from ndcube import NDCube
from prefect import flow, get_run_logger

from punchbowl.data import NormalizedMetadata
from punchbowl.data.meta import set_spacecraft_location_to_earth
from punchbowl.level3.f_corona_model import subtract_f_corona_background_task
from punchbowl.level3.low_noise import create_low_noise_task
from punchbowl.level3.polarization import convert_polarization
Expand All @@ -14,7 +16,8 @@
def level3_PIM_flow(data_list: list[str] | list[NDCube], # noqa: N802
before_f_corona_model_path: str,
after_f_corona_model_path: str,
output_filename: str | None = None) -> list[NDCube]:
output_filename: str | None = None,
reference_time: datetime | None = None) -> list[NDCube]: # noqa: ARG001
"""Level 3 PIM flow."""
logger = get_run_logger()

Expand All @@ -24,9 +27,15 @@ def level3_PIM_flow(data_list: list[str] | list[NDCube], # noqa: N802
before_f_corona_model_path,
after_f_corona_model_path) for d in data_list]
output_meta = NormalizedMetadata.load_template("PIM", "3")

out_list = [NDCube(data=d.data, wcs=d.wcs, meta=output_meta) for d in data_list]
for o, d in zip(out_list, data_list, strict=True):
for o, d in zip(out_list, data_list, strict=True):
o.meta["DATE"] = datetime.now().isoformat()
o.meta["DATE-AVG"] = d.meta["DATE-AVG"].value
o.meta["DATE-OBS"] = d.meta["DATE-OBS"].value
o.meta["DATE-BEG"] = d.meta["DATE-BEG"].value
o.meta["DATE-END"] = d.meta["DATE-END"].value
o = set_spacecraft_location_to_earth(o) # noqa: PLW2901

logger.info("ending level 3 PIM flow")

Expand Down Expand Up @@ -54,6 +63,7 @@ def level3_core_flow(data_list: list[str] | list[NDCube],
data_list = [convert_polarization(d) for d in data_list]
logger.info("ending level 3 core flow")


if output_filename is not None:
output_image_task(data_list[0], output_filename)

Expand Down
9 changes: 8 additions & 1 deletion punchbowl/level3/polarization.py
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
from datetime import datetime

import astropy.units as u
import numpy as np
from ndcube import NDCollection, NDCube
from prefect import get_run_logger
from solpolpy import resolve

from punchbowl.data.meta import NormalizedMetadata
from punchbowl.data.meta import NormalizedMetadata, set_spacecraft_location_to_earth
from punchbowl.prefect import punch_task


Expand All @@ -30,8 +32,13 @@ def convert_polarization(
new_wcs = input_data.wcs.copy()

output_meta = NormalizedMetadata.load_template("PTM", "3")
output_meta["DATE"] = datetime.now().isoformat()
output_meta["DATE-AVG"] = input_data.meta["DATE-AVG"].value
output_meta["DATE-OBS"] = input_data.meta["DATE-OBS"].value
output_meta["DATE-BEG"] = input_data.meta["DATE-BEG"].value
output_meta["DATE-END"] = input_data.meta["DATE-END"].value
output = NDCube(data=new_data, wcs=new_wcs, meta=output_meta)
output = set_spacecraft_location_to_earth(output)

logger.info("convert2bpb finished")
output.meta.history.add_now("LEVEL3-convert2bpb", "Convert MZP to BpB")
Expand Down
Loading

0 comments on commit 3b3b496

Please sign in to comment.