Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix qlook function for file saving #144

Merged
merged 3 commits into from
Dec 5, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 71 additions & 30 deletions decode/qlook.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@
DEFAULT_FORMAT = "png"
DEFAULT_FREQUENCY_UNITS = "GHz"
DEFAULT_INCL_MKID_IDS = None
DEFAULT_OUTDIR = Path()
DEFAULT_OVERWRITE = False
DEFAULT_SKYCOORD_GRID = "6 arcsec"
DEFAULT_SKYCOORD_UNITS = "arcsec"
SIGMA_OVER_MAD = 1.4826
Expand All @@ -43,7 +45,8 @@ def pswsc(
data_type: Literal["df/f", "brightness", None] = DEFAULT_DATA_TYPE,
frequency_units: str = DEFAULT_FREQUENCY_UNITS,
format: str = DEFAULT_FORMAT,
outdir: Path = Path(),
outdir: Path = DEFAULT_OUTDIR,
overwrite: bool = DEFAULT_OVERWRITE,
) -> Path:
"""Quick-look at a PSW observation with sky chopper.

Expand All @@ -58,6 +61,7 @@ def pswsc(
frequency_units: Units of the frequency axis.
format: Output data format of the quick-look result.
outdir: Output directory for the quick-look result.
overwrite: Whether to overwrite the output if it exists.

Returns:
Absolute path of the saved file.
Expand All @@ -77,10 +81,11 @@ def pswsc(
spec = da_sub.mean("scan")

# save result
filename = Path(dems).with_suffix(f".pswsc.{format}").name
file_name = Path(dems).with_suffix(f".pswsc.{format}").name
file_path = Path(outdir) / file_name

if format in DATA_FORMATS:
return save_qlook(spec, Path(outdir) / filename)
return save_qlook(spec, file_path, overwrite=overwrite)

fig, axes = plt.subplots(1, 2, figsize=DEFAULT_FIGSIZE)

Expand All @@ -96,7 +101,7 @@ def pswsc(
ax.grid(True)

fig.tight_layout()
return save_qlook(fig, Path(outdir) / filename)
return save_qlook(fig, file_path, overwrite=overwrite)


def raster(
Expand All @@ -111,7 +116,8 @@ def raster(
skycoord_grid: str = DEFAULT_SKYCOORD_GRID,
skycoord_units: str = DEFAULT_SKYCOORD_UNITS,
format: str = DEFAULT_FORMAT,
outdir: Path = Path(),
outdir: Path = DEFAULT_OUTDIR,
overwrite: bool = DEFAULT_OVERWRITE,
) -> Path:
"""Quick-look at a raster scan observation.

Expand All @@ -134,6 +140,7 @@ def raster(
skycoord_units: Units of the sky coordinate axes.
format: Output image format of quick-look result.
outdir: Output directory for the quick-look result.
overwrite: Whether to overwrite the output if it exists.

Returns:
Absolute path of the saved file.
Expand Down Expand Up @@ -174,10 +181,11 @@ def raster(
cont = cube.weighted(weight.fillna(0)).mean("chan")

# save result
filename = Path(dems).with_suffix(f".raster.{format}").name
file_name = Path(dems).with_suffix(f".raster.{format}").name
file_path = Path(outdir) / file_name

if format in DATA_FORMATS:
return save_qlook(cont, Path(outdir) / filename)
return save_qlook(cont, file_path, overwrite=overwrite)

fig, axes = plt.subplots(1, 2, figsize=(12, 5.5))

Expand All @@ -202,7 +210,7 @@ def raster(
ax.grid(True)

fig.tight_layout()
return save_qlook(fig, Path(outdir) / filename)
return save_qlook(fig, file_path, overwrite=overwrite)


def skydip(
Expand All @@ -215,7 +223,8 @@ def skydip(
chan_weight: Literal["uniform", "std", "std/tx"] = "std/tx",
pwv: Literal["0.5", "1.0", "2.0", "3.0", "4.0", "5.0"] = "5.0",
format: str = DEFAULT_FORMAT,
outdir: Path = Path(),
outdir: Path = DEFAULT_OUTDIR,
overwrite: bool = DEFAULT_OVERWRITE,
) -> Path:
"""Quick-look at a skydip observation.

Expand All @@ -236,6 +245,7 @@ def skydip(
the atmospheric transmission when chan_weight is std/tx.
format: Output image format of quick-look result.
outdir: Output directory for the quick-look result.
overwrite: Whether to overwrite the output if it exists.

Returns:
Absolute path of the saved file.
Expand All @@ -255,10 +265,11 @@ def skydip(
series = da_on.weighted(weight.fillna(0)).mean("chan")

# save result
filename = Path(dems).with_suffix(f".skydip.{format}").name
file_name = Path(dems).with_suffix(f".raster.{format}").name
file_path = Path(outdir) / file_name

if format in DATA_FORMATS:
return save_qlook(series, Path(outdir) / filename)
return save_qlook(series, file_path, overwrite=overwrite)

fig, axes = plt.subplots(1, 2, figsize=DEFAULT_FIGSIZE)

Expand All @@ -273,7 +284,7 @@ def skydip(
ax.grid(True)

fig.tight_layout()
return save_qlook(fig, Path(outdir) / filename)
return save_qlook(series, file_path, overwrite=overwrite)


def still(
Expand All @@ -286,7 +297,8 @@ def still(
chan_weight: Literal["uniform", "std", "std/tx"] = "std/tx",
pwv: Literal["0.5", "1.0", "2.0", "3.0", "4.0", "5.0"] = "5.0",
format: str = DEFAULT_FORMAT,
outdir: Path = Path(),
outdir: Path = DEFAULT_OUTDIR,
overwrite: bool = DEFAULT_OVERWRITE,
) -> Path:
"""Quick-look at a still observation.

Expand All @@ -307,6 +319,7 @@ def still(
the atmospheric transmission when chan_weight is std/tx.
format: Output data format of the quick-look result.
outdir: Output directory for the quick-look result.
overwrite: Whether to overwrite the output if it exists.

Returns:
Absolute path of the saved file.
Expand All @@ -325,10 +338,11 @@ def still(
series = da.weighted(weight.fillna(0)).mean("chan")

# save result
filename = Path(dems).with_suffix(f".still.{format}").name
file_name = Path(dems).with_suffix(f".still.{format}").name
file_path = Path(outdir) / file_name

if format in DATA_FORMATS:
return save_qlook(series, Path(outdir) / filename)
return save_qlook(series, file_path, overwrite=overwrite)

fig, axes = plt.subplots(1, 2, figsize=DEFAULT_FIGSIZE)

Expand All @@ -343,7 +357,7 @@ def still(
ax.grid(True)

fig.tight_layout()
return save_qlook(fig, Path(outdir) / filename)
return save_qlook(series, file_path, overwrite=overwrite)


def zscan(
Expand All @@ -356,7 +370,8 @@ def zscan(
chan_weight: Literal["uniform", "std", "std/tx"] = "std/tx",
pwv: Literal["0.5", "1.0", "2.0", "3.0", "4.0", "5.0"] = "5.0",
format: str = DEFAULT_FORMAT,
outdir: Path = Path(),
outdir: Path = DEFAULT_OUTDIR,
overwrite: bool = DEFAULT_OVERWRITE,
) -> Path:
"""Quick-look at an observation of subref axial focus scan.

Expand All @@ -377,6 +392,7 @@ def zscan(
the atmospheric transmission when chan_weight is std/tx.
format: Output image format of quick-look result.
outdir: Output directory for the quick-look result.
overwrite: Whether to overwrite the output if it exists.

Returns:
Absolute path of the saved file.
Expand All @@ -396,10 +412,11 @@ def zscan(
series = da_on.weighted(weight.fillna(0)).mean("chan")

# save result
filename = Path(dems).with_suffix(f".zscan.{format}").name
file_name = Path(dems).with_suffix(f".zscan.{format}").name
file_path = Path(outdir) / file_name

if format in DATA_FORMATS:
return save_qlook(series, Path(outdir) / filename)
return save_qlook(series, file_path, overwrite=overwrite)

fig, axes = plt.subplots(1, 2, figsize=DEFAULT_FIGSIZE)

Expand All @@ -414,7 +431,7 @@ def zscan(
ax.grid(True)

fig.tight_layout()
return save_qlook(fig, Path(outdir) / filename)
return save_qlook(series, file_path, overwrite=overwrite)


def mean_in_time(dems: xr.DataArray) -> xr.DataArray:
Expand Down Expand Up @@ -565,28 +582,52 @@ def load_dems(
raise ValueError("Data type could not be inferred.")


def save_qlook(qlook: Union[Figure, xr.DataArray], filename: Path) -> Path:
def save_qlook(
qlook: Union[Figure, xr.DataArray],
file: Path,
/,
*,
overwrite: bool = False,
) -> Path:
"""Save a quick look result to a file with given format.

Args:
qlook: Matplotlib figure or DataArray to be saved.
filename: Path of the saved file.
file: Path of the saved file.
overwrite: Whether to overwrite the file if it exists.

Returns:
Absolute path of the saved file.

"""
path = Path(file).expanduser().resolve()

if path.exists() and not overwrite:
raise FileExistsError(f"{path} already exists.")

if isinstance(qlook, Figure):
qlook.savefig(filename)
elif (ext := "".join(filename.suffixes)) == ".csv":
qlook.savefig(path)
return path

if path.name.endswith(".csv"):
name = qlook.attrs["data_type"]
qlook.to_dataset(name=name).to_pandas().to_csv(filename)
elif ext == ".nc":
qlook.to_netcdf(filename)
elif ext == ".zarr" or format == ".zarr.zip":
qlook.to_zarr(filename, mode="w")
ds = qlook.to_dataset(name=name)
ds.to_pandas().to_csv(path)
return path

if path.name.endswith(".nc"):
qlook.to_netcdf(path)
return path

if path.name.endswith(".zarr"):
qlook.to_zarr(path, mode="w")
return path

if path.name.endswith(".zarr.zip"):
qlook.to_zarr(path, mode="w")
return path

return Path(filename).expanduser().resolve()
raise ValueError("Extension of filename is not valid.")


def main() -> None:
Expand Down