diff --git a/hera_notebook_templates/notebooks/lststack.ipynb b/hera_notebook_templates/notebooks/lststack.ipynb index 494e2d8..7aa829e 100644 --- a/hera_notebook_templates/notebooks/lststack.ipynb +++ b/hera_notebook_templates/notebooks/lststack.ipynb @@ -24,7 +24,7 @@ "id": "543efff3", "metadata": {}, "source": [ - "**by Steven Murray**, last updated 27th Mar, 2024." + "**by Steven Murray, Tyler Cox and Josh Dillon**, last updated 4th Jun, 2024." ] }, { @@ -80,12 +80,20 @@ "import h5py\n", "import attrs\n", "\n", + "from scipy import linalg, constants, signal\n", + "from hera_filters import dspec\n", + "\n", "from pyuvdata import UVData\n", + "from hera_cal.abscal import complex_phase_abscal\n", "from hera_cal import lst_stack as lstbin\n", "from hera_cal.lst_stack.config import LSTConfig\n", "from hera_cal.red_groups import RedundantGroups\n", "from hera_cal.lst_stack import metrics as lstmet\n", "from hera_cal.lst_stack import stats as lststat\n", + "from hera_cal.lst_stack.calibration import lstbin_absolute_calibration\n", + "from hera_cal.lst_stack import averaging as avg\n", + "from hera_cal import vis_clean, redcal, abscal, utils\n", + "from hera_qm.time_series_metrics import true_stretches\n", "import importlib" ] }, @@ -120,8 +128,9 @@ }, "outputs": [], "source": [ - "fileconf: str = \"/lustre/aoc/projects/hera/h6c-analysis/IDR2/lstbin-outputs/redavg-smoothcal-notebook/file-config.h5\"\n", + "fileconf: str = \"/lustre/aoc/projects/hera/h6c-analysis/IDR2/lstbin-outputs/redavg-smoothcal-inpaint/file-config.h5\"\n", "fileidx: int = 380\n", + "blchunk: int = 0\n", "\n", "papermill_input_path: str = \"\"\n", "papermill_output_path: str = \"\"\n", @@ -131,13 +140,14 @@ "save_metric_data: bool = True\n", "plot_n_worst: int = 5\n", "do_simultaneous_inpainting: bool = True\n", - "do_extra_flagging: bool = True\n", + "do_extra_flagging: bool = False\n", + "do_lstcal: bool = True\n", "\n", "outdir: str = \".\"\n", "bl_chunk_size: int = 0\n", "rephase: bool = True\n", "vis_units: str = \"Jy\"\n", - "fname_format: str = '{inpaint_mode}/zen.{kind}.{lst:7.5f}.sum.uvh5'\n", + "fname_format: str = '{inpaint_mode}/zen.{kind}.{lst:7.5f}.{blchunk}.sum.uvh5'\n", "overwrite: bool = True\n", "write_med_mad: bool = False\n", "do_inpainted_average: bool = False\n", @@ -147,13 +157,17 @@ "plot_every: int = 1\n", "\n", "# In-painting config\n", - "delay_filter_horizon: float = 1.0\n", - "delay_filter_standoff: float = 0.0 # ns\n", - "delay_filter_mindelay: float = 150.0 # ns\n", - "delay_filter_eigencutoff: float = 1e-12\n", - "inpaint_regularization: float = 1e-5 # reasonable values are between 1e-2 and 1e-5\n", - "inpaint_mindelay: float = 500.0 # ns\n", - "inpaint_cache_dir: str = \"\" # leave empty to NOT write cache to file. If a file, all LST-bins should use the same file so they can take advantage of it.\n", + "inpaint_horizon: float = 1.0\n", + "inpaint_standoff: float = 0.0 # ns\n", + "inpaint_eigencutoff: float = 1e-12\n", + "inpaint_mindelay: float = 500.0 # ns\n", + "inpaint_max_gap_factor: float = 2.0\n", + "inpaint_max_convolved_flag_frac: float = 0.667\n", + "inpaint_sample_cov_fraction: float = 0.0 # Default zero uses variance, one uses full sample covariance\n", + "inpaint_use_unbiased_estimator: bool = False # Default False slight over estimate of the covariance, but guaranteed to be non-negative\n", + "\n", + "FM_low_freq: float = 87.5 # in MHz\n", + "FM_high_freq: float = 108.0 # in MHz\n", "\n", "# Flagging Configuration\n", "zscore_threshold: float = 5 # Value of |Z| above which data are flagged. \n", @@ -162,51 +176,6 @@ "max_flagging_iterations: int = 15 # Maximum number of iterations to perform when flagging." ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "1f777f41-737c-441d-a7dc-6f02ac402959", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "# TODO: using the filter cache doesn't work yet, because the keys are tuples, not strings\n", - "class FilterCache(UserDict):\n", - " def __init__(self, cachedir: str | Path | None = None, *args, **kwargs):\n", - " if cachedir:\n", - " self._cachedir = Path(cachedir)\n", - " else:\n", - " self._cachedir = None\n", - " \n", - " super().__init__(*args,**kwargs)\n", - " \n", - " def __getitem__(self, item):\n", - " try:\n", - " return super().__getitem__(item)\n", - " except KeyError as e:\n", - " if self._cachedir is None:\n", - " raise\n", - " \n", - " possible_keys = [pth.name for pth in self._cachedir.glob(\"*\")]\n", - " \n", - " if item in possible_keys:\n", - " with (self._cachedir / item).open('rb') as fl:\n", - " out = pickle.load(fl)\n", - " self[item] = out\n", - " return out\n", - " \n", - " raise\n", - " \n", - " def __setitem__(self, item, value):\n", - " if self._cachedir is not None:\n", - " with (self._cachedir / item).open('wb') as fl:\n", - " pickle.dump(fl, value)\n", - " \n", - " super().__setitem__(item, value)\n", - " " - ] - }, { "cell_type": "code", "execution_count": null, @@ -242,7 +211,9 @@ "cell_type": "code", "execution_count": null, "id": "4117ac48-9b06-4362-b6e5-954fe4e0ac89", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "if make_plots or save_metric_data:\n", @@ -386,7 +357,9 @@ "cell_type": "code", "execution_count": null, "id": "adf225cc-f865-4945-9986-83bd9375938b", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "if do_simultaneous_inpainting and do_inpainted_average:\n", @@ -412,7 +385,9 @@ "else:\n", " bl_chunk_size = min(bl_chunk_size, len(stackconf.antpairs))\n", "\n", - "n_bl_chunks = int(np.ceil(len(stackconf.antpairs) / bl_chunk_size))" + "n_bl_chunks = int(np.ceil(len(stackconf.antpairs) / bl_chunk_size))\n", + "\n", + "print(f\"This notebook is processing chunk {blchunk+1} of {n_bl_chunks} baseline chunks, each with {bl_chunk_size} baselines.\")" ] }, { @@ -430,18 +405,6 @@ ")" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "37c0176d-1e50-45a6-a8c7-07b1f4737cec", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "_INPAINT_CACHE_ = FilterCache(inpaint_cache_dir)" - ] - }, { "cell_type": "code", "execution_count": null, @@ -482,7 +445,9 @@ "cell_type": "code", "execution_count": null, "id": "378e26d9", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "if len(set(sum((x.tolist() for x in stackconf.time_indices), start=[]))) != stackconf.n_lsts:\n", @@ -495,7 +460,9 @@ "cell_type": "code", "execution_count": null, "id": "a6530710", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "# If this notebook is making plots, and is being run through PAPERMILL, output an empty\n", @@ -506,6 +473,19 @@ " pth.touch()" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "f6c956a0-4fa8-49a9-b866-ac1f2ec29e9d", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "data_is_redundantly_averaged = stackconf.config.is_redundantly_averaged\n", + "print(f\"This data {'is' if data_is_redundantly_averaged else 'is not'} redundantly averaged.\")" + ] + }, { "cell_type": "markdown", "id": "8dcb2696", @@ -709,70 +689,141 @@ "}" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "84d270e0-fc06-475f-9f3b-9b939b6346fe", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "inpaint_bands = [(0, FM_low_freq), (FM_high_freq, np.inf)] # default below and above FM\n", + "\n", + "# Get slices for the inpaint bands\n", + "_inp = []\n", + "for _bnd in inpaint_bands:\n", + " idx = np.nonzero((stackconf.config.datameta.freq_array >= _bnd[0] * 1e6) & (stackconf.config.datameta.freq_array < _bnd[1]*1e6))[0]\n", + " _inp.append(slice(idx[0], idx[-1] + 1))\n", + "inpaint_bands = _inp\n", + "print(\"Using the following bands for inpainting (channels):\")\n", + "for bnd in inpaint_bands:\n", + " print(bnd)" + ] + }, { "cell_type": "markdown", - "id": "af3b55b2", - "metadata": {}, + "id": "361a1e4d-42c6-4800-824f-1e190d114e2f", + "metadata": { + "execution": { + "iopub.execute_input": "2024-06-04T18:20:46.071514Z", + "iopub.status.busy": "2024-06-04T18:20:46.070413Z", + "iopub.status.idle": "2024-06-04T18:20:46.084305Z", + "shell.execute_reply": "2024-06-04T18:20:46.081600Z", + "shell.execute_reply.started": "2024-06-04T18:20:46.071441Z" + } + }, "source": [ - "## LST-bin the Autos" + "## Perform Initial Stacking of Autos and Crosses" ] }, { "cell_type": "code", "execution_count": null, - "id": "f54691c7", + "id": "759dfcac-3eb9-41d6-a7ed-ca640bbb9047", "metadata": { "tags": [] }, "outputs": [], "source": [ - "def make_auto_plot(auto_stacks: list[UVData], lstbin: list[dict]):\n", - " \n", - " fig, ax = plt.subplots(\n", - " len(stackconf.autopairs)*len(stackconf.pols), len(auto_stacks), \n", - " sharex=True, sharey=True, squeeze=False, constrained_layout=True,\n", - " figsize=(12, 6)\n", - " )\n", - "\n", - " for i, (stack, avg) in enumerate(zip(auto_stacks, lstbin)):\n", - " for j, autopair in enumerate(stackconf.autopairs):\n", - " for p, pol in enumerate(stackconf.pols):\n", - " axx = ax[j*len(stackconf.pols) + p, i]\n", - " \n", - " for k, t in enumerate(stack.time_array[::stack.Nbls]):\n", - " flg = stack.get_flags(autopair + (pol,))[k]\n", - " d = stack.get_data(autopair+(pol,))[k]\n", - " \n", - " axx.plot(\n", - " stack.freq_array / 1e6,\n", - " np.where(flg, np.nan, d.real),\n", - " label=f\"{int(t)}\" if not p else None,\n", - " **styles[int(t)]\n", - " )\n", - " axx.set_yscale('log')\n", - " axx.set_title(f\"Pair {autopair}, pol={pol}, LST {stackconf.lst_grid[i]*12/np.pi:.3f} hr\")\n", - "\n", - " # plot the mean\n", - " axx.plot(\n", - " stack.freq_array / 1e6,\n", - " np.where(avg['flags'][j, :, p], np.nan, avg['data'][j, :, p].real),\n", - " label='LSTBIN',\n", - " color='k', lw=2\n", - " )\n", - " \n", - " ax[0,0].legend(ncols=3)" + "auto_stacks, autos_lstavg = stack_blchunk('autos') # Auto-stacks\n", + "cross_stacks, cross_lstavg = stack_blchunk(blchunk) # Cross-stacks" + ] + }, + { + "cell_type": "markdown", + "id": "0e1501b7-2715-4821-8775-0a9b87e38fd5", + "metadata": {}, + "source": [ + "## LST-Bin Calibration" ] }, { "cell_type": "code", "execution_count": null, - "id": "29c9e7b2", + "id": "08c855c9-0869-4d91-bd45-9dd7e21a6599", "metadata": { "tags": [] }, "outputs": [], "source": [ - "auto_stacks, autos_lstavg = stack_blchunk('autos')" + "%%time\n", + "if do_lstcal and n_bl_chunks==1: # can't fit all the baselines in memory if not redavg'd, and need all of them at once to do lstcal\n", + " for i, (stack, lstavg_model, auto_model) in enumerate(zip(cross_stacks, cross_lstavg, autos_lstavg)):\n", + " calibration_parameters, gains = lstbin_absolute_calibration(\n", + " stack=stack, \n", + " model=lstavg_model['data'], \n", + " all_reds=reds_with_pols, \n", + " inpaint_bands=inpaint_bands,\n", + " auto_stack=auto_stacks[i],\n", + " auto_model=auto_model['data'],\n", + " calibrate_inplace=True, # calibrate inplace\n", + " run_phase_cal=True, # run phase calibration\n", + " run_amplitude_cal=True, # run amplitude calibration\n", + " )\n", + " \n", + " # Write out calibration parameters and metadata\n", + " calibration_parameters[\"freqs\"] = stack.freq_array\n", + " calibration_parameters[\"flags\"] = stack.flags[:, 0, :]\n", + " calibration_parameters[\"times\"] = stack.times\n", + " calibration_parameters[\"antpairs\"] = stack.antpairs\n", + " calibration_parameters[\"lst\"] = stackconf.lst_grid[i]\n", + " \n", + " # Get the calibration filename\n", + " cal_fname = lstbin.io.format_outfile_name(\n", + " lst=stackconf.lst_grid_edges[i],\n", + " pols=stackconf.pols,\n", + " fname_format=fname_format.replace(\"sum.uvh5\", \"lstcal.npz\"),\n", + " inpaint_mode=True,\n", + " lst_branch_cut=stackconf.properties[\"lst_branch_cut\"],\n", + " kind=\"LST\"\n", + " )\n", + " outfile = outdir / cal_fname\n", + " \n", + " np.savez(outfile, **calibration_parameters)\n", + " \n", + " \n", + " # Recompute auto and cross averages after calibration - needed for STD files \n", + " cross_lstavg = [\n", + " lstbin.averaging.reduce_lst_bins(\n", + " lststack=stack,\n", + " inpainted_mode=False,\n", + " get_mad=write_med_mad,\n", + " ) for stack in cross_stacks\n", + " ]\n", + " autos_lstavg = [\n", + " lstbin.averaging.reduce_lst_bins(\n", + " lststack=stack,\n", + " inpainted_mode=False,\n", + " get_mad=write_med_mad,\n", + " ) for stack in auto_stacks\n", + " ]" + ] + }, + { + "cell_type": "markdown", + "id": "af3b55b2", + "metadata": {}, + "source": [ + "## Autos" + ] + }, + { + "cell_type": "markdown", + "id": "c646623f", + "metadata": {}, + "source": [ + "### Compute Stats for Autos" ] }, { @@ -794,34 +845,28 @@ "id": "ea584858", "metadata": {}, "source": [ - "### Inpaint" + "### Inpaint Autos" ] }, { "cell_type": "code", "execution_count": null, - "id": "20fc85fe", - "metadata": {}, + "id": "f62b3534", + "metadata": { + "tags": [] + }, "outputs": [], "source": [ - "inpaint_bands = [(0, 87.5), (108.0, 250.0)] # default below and above FM\n", - "\n", - "# Get slices for the inpaint bands\n", - "_inp = []\n", - "for _bnd in inpaint_bands:\n", - " idx = np.nonzero((stackconf.config.datameta.freq_array >= _bnd[0] * 1e6) & (stackconf.config.datameta.freq_array < _bnd[1]*1e6))[0]\n", - " _inp.append(slice(idx[0], idx[-1] + 1))\n", - "inpaint_bands = _inp\n", - "print(\"Using the following bands for inpainting (channels):\")\n", - "for bnd in inpaint_bands:\n", - " print(bnd)" + "_INPAINT_CACHE_ = {}" ] }, { "cell_type": "code", "execution_count": null, "id": "1da6906c", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "if do_simultaneous_inpainting:\n", @@ -829,18 +874,22 @@ "\n", " for i, stack in enumerate(auto_stacks):\n", " \n", - " _avg, dpss_models = lstbin.averaging.average_and_inpaint_simultaneously(\n", + " _avg, dpss_models = avg.average_and_inpaint_simultaneously(\n", + " stack,\n", " stack,\n", " inpaint_bands = inpaint_bands,\n", " return_models = make_plots,\n", " cache = _INPAINT_CACHE_,\n", " filter_properties = {\n", " \"min_dly\": inpaint_mindelay, \n", - " \"horizon\": delay_filter_horizon,\n", - " \"standoff\": delay_filter_standoff, \n", + " \"horizon\": inpaint_horizon,\n", + " \"standoff\": inpaint_standoff, \n", " },\n", - " eigenval_cutoff=[delay_filter_eigencutoff], \n", - " suppression_factors=[delay_filter_eigencutoff], \n", + " eigenval_cutoff=[inpaint_eigencutoff], \n", + " max_gap_factor=inpaint_max_gap_factor,\n", + " max_convolved_flag_frac=inpaint_max_convolved_flag_frac,\n", + " use_unbiased_estimator=inpaint_use_unbiased_estimator,\n", + " sample_cov_fraction=inpaint_sample_cov_fraction,\n", " )\n", "\n", " auto_inpaint_dpss_models.append(dpss_models)\n", @@ -859,34 +908,68 @@ { "cell_type": "code", "execution_count": null, - "id": "6a34526c", + "id": "1fccd669-53d7-4888-8382-831e028966a3", "metadata": { "tags": [] }, "outputs": [], "source": [ - "if make_plots:\n", - " make_auto_plot(auto_stacks, autos_lstavg);" - ] - }, - { - "cell_type": "markdown", - "id": "ae48cd95", - "metadata": {}, - "source": [ - "## Cross-Pairs" + "def make_auto_plot(auto_stacks: list[UVData], lstbin: list[dict]):\n", + " \n", + " fig, ax = plt.subplots(\n", + " len(stackconf.autopairs)*len(stackconf.pols), len(auto_stacks), \n", + " sharex=True, sharey=True, squeeze=False, constrained_layout=True,\n", + " figsize=(12, 6)\n", + " )\n", + "\n", + " for i, (stack, avg) in enumerate(zip(auto_stacks, lstbin)):\n", + " for j, autopair in enumerate(stackconf.autopairs):\n", + " for p, pol in enumerate(stackconf.pols):\n", + " axx = ax[j*len(stackconf.pols) + p, i]\n", + " \n", + " for k, t in enumerate(stack.time_array[::stack.Nbls]):\n", + " flg = stack.get_flags(autopair + (pol,))[k]\n", + " d = stack.get_data(autopair+(pol,))[k]\n", + " \n", + " axx.plot(\n", + " stack.freq_array / 1e6,\n", + " np.where(flg, np.nan, d.real),\n", + " label=f\"{int(t)}\" if not p else None,\n", + " **styles[int(t)]\n", + " )\n", + " axx.set_yscale('log')\n", + " axx.set_title(f\"Pair {autopair}, pol={pol}, LST {stackconf.lst_grid[i]*12/np.pi:.3f} hr\")\n", + "\n", + " # plot the mean\n", + " axx.plot(\n", + " stack.freq_array / 1e6,\n", + " np.where(avg['flags'][j, :, p], np.nan, avg['data'][j, :, p].real),\n", + " label='LSTBIN',\n", + " color='k', lw=2\n", + " )\n", + " \n", + " ax[0,0].legend(ncols=3)" ] }, { "cell_type": "code", "execution_count": null, - "id": "2994dd78-9fef-4006-a2d3-3ed7ae02e353", + "id": "6a34526c", "metadata": { "tags": [] }, "outputs": [], "source": [ - "cross_stacks, cross_lstavg = stack_blchunk(0)" + "if make_plots:\n", + " make_auto_plot(auto_stacks, autos_lstavg)" + ] + }, + { + "cell_type": "markdown", + "id": "ae48cd95", + "metadata": {}, + "source": [ + "## Cross-Pairs" ] }, { @@ -989,7 +1072,9 @@ "cell_type": "code", "execution_count": null, "id": "c2f51883-9596-4fc1-992e-847d719e9719", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "zsquare = [lstmet.get_squared_zscores_flagged(stack, auto_stats=stats) for stack, stats in zip(cross_stacks, auto_stats)]" @@ -1016,7 +1101,7 @@ "id": "ca98fc9f-a22e-4a2d-8ff1-8bccfdb453f5", "metadata": {}, "source": [ - "### Inpaint" + "### Inpaint Crosses" ] }, { @@ -1024,25 +1109,7 @@ "id": "461a78dc-d227-4479-8ebb-546e6ac5568e", "metadata": {}, "source": [ - "We simultaneously inpaint and average the data with the flags we're given. We don't try and inpaint the stack itself (i.e. on a nightly basis), since we have only one solution per LST bin (per freq and baseline), and there's no useful information gained by doing so (one might think that the newly-inpainted nightly solution might be useful in terms of getting a new $Z^2$ score for the statistics, but the $Z^2$ for the inpainted data is not well-defined: it has zero nsamples). " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d33bc2e2-c610-4046-96f6-30d5585d9f17", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "cross_lstavg = [\n", - " lstbin.averaging.reduce_lst_bins(\n", - " lststack=stack,\n", - " inpainted_mode=False,\n", - " get_mad=write_med_mad,\n", - " ) for stack in cross_stacks\n", - "]" + "We simultaneously inpaint and average the data with the flags we're given, using the covariance of a DPSS fit to the mean as a constraint." ] }, { @@ -1068,29 +1135,31 @@ }, "outputs": [], "source": [ + "%%time\n", "if do_simultaneous_inpainting:\n", " inpaint_dpss_models = []\n", "\n", - " for i, stack in enumerate(cross_stacks):\n", - " \n", - " _avg, dpss_models = lstbin.averaging.average_and_inpaint_simultaneously(\n", + " for i, (stack, auto_stack) in enumerate(zip(cross_stacks, auto_stacks)):\n", + " _avg, dpss_models = avg.average_and_inpaint_simultaneously(\n", " stack,\n", + " auto_stack,\n", " inpaint_bands = inpaint_bands,\n", " return_models = make_plots,\n", " cache = _INPAINT_CACHE_,\n", " filter_properties = {\n", " \"min_dly\": inpaint_mindelay, \n", - " \"horizon\": delay_filter_horizon,\n", - " \"standoff\": delay_filter_standoff, \n", + " \"horizon\": inpaint_horizon,\n", + " \"standoff\": inpaint_standoff, \n", " },\n", - " eigenval_cutoff=[delay_filter_eigencutoff], \n", - " suppression_factors=[delay_filter_eigencutoff], \n", + " eigenval_cutoff=[inpaint_eigencutoff], \n", + " max_gap_factor=inpaint_max_gap_factor,\n", + " max_convolved_flag_frac=inpaint_max_convolved_flag_frac,\n", + " use_unbiased_estimator=inpaint_use_unbiased_estimator,\n", + " sample_cov_fraction=inpaint_sample_cov_fraction,\n", " )\n", - "\n", " inpaint_dpss_models.append(dpss_models)\n", " cross_lstavg[i]['data'] = _avg['data']\n", - " cross_lstavg[i]['flags'] = _avg['flags']\n", - "\n" + " cross_lstavg[i]['flags'] = _avg['flags']\n" ] }, { @@ -1188,7 +1257,7 @@ " )\n", " ax[kk].plot(fq, np.abs(cross_lstavg[0]['data'][apidx, band, polidx][subband]), lw=3, ls='--', color='red', label='simul. inpaint')\n", "\n", - " ax[kk].plot(fq, np.abs(inpaint_dpss_models[0][blpol][band][subband]), lw=2, color='purple', ls=':', label='model')\n", + " #ax[kk].plot(fq, np.abs(inpaint_dpss_models[0][blpol][band][subband]), lw=2, color='purple', ls=':', label='model')\n", "\n", " gotlabel = False\n", " for i, jd in enumerate(cross_stacks[0].nights):\n", @@ -1237,12 +1306,12 @@ "source": [ "if do_inpainted_average:\n", " cross_lstavg = [\n", - " lstbin.averaging.reduce_lst_bins(\n", - " lststack=stack,\n", - " inpainted_mode=True,\n", - " get_mad=write_med_mad,\n", - " ) for stack in cross_stacks\n", - "]" + " lstbin.averaging.reduce_lst_bins(\n", + " lststack=stack,\n", + " inpainted_mode=True,\n", + " get_mad=write_med_mad,\n", + " ) for stack in cross_stacks\n", + " ]" ] }, { @@ -1255,11 +1324,11 @@ "outputs": [], "source": [ "# make it a bit easier to create the outfiles\n", - "uvd_template = lstbin.io.create_empty_uvd(\n", + "create_empty_uvd = partial(\n", + " lstbin.io.create_empty_uvd,\n", " pols=stackconf.pols,\n", " file_list=stackconf.matched_metas,\n", " history=history,\n", - " antpairs=stackconf.autopairs + stackconf.antpairs,\n", " start_jd=stackconf.properties['first_jd'],\n", " freq_min=freq_min,\n", " freq_max=freq_max,\n", @@ -1270,9 +1339,18 @@ "\n", "create_file = partial(\n", " lstbin.io.create_lstbin_output_file,\n", - " uvd_template=uvd_template,\n", + "# uvd_template=uvd_template,\n", " outdir=outdir,\n", " overwrite=overwrite,\n", + ")\n", + "\n", + "get_fname = partial(\n", + " lstbin.io.format_outfile_name,\n", + " lst=stackconf.lst_grid_edges[0],\n", + " pols=stackconf.pols,\n", + " #fname_format=fname_format,\n", + " inpaint_mode=True,\n", + " lst_branch_cut=stackconf.properties[\"lst_branch_cut\"],\n", ")" ] }, @@ -1290,18 +1368,21 @@ " kinds = [\"LST\", \"STD\"]\n", " if write_med_mad:\n", " kinds += [\"MED\", \"MAD\"]\n", + " \n", " for kind in kinds:\n", + " if blchunk == 0:\n", + " # Save the autos as well.\n", + " fname = get_fname(\n", + " fname_format = fname_format.replace(\"{blchunk}\", \"autos\"), kind=kind\n", + " )\n", + " auto_uvd_template = create_empty_uvd(antpairs=auto_stack.antpairs)\n", + " out_files[(kind, 'autos')] = create_file(fname=fname, uvd_template=auto_uvd_template)\n", "\n", - " fname = lstbin.io.format_outfile_name(\n", - " lst=stackconf.lst_grid_edges[0],\n", - " pols=stackconf.pols,\n", - " fname_format=fname_format,\n", - " inpaint_mode=True,\n", - " lst_branch_cut=stackconf.properties[\"lst_branch_cut\"],\n", - " kind=kind\n", + " fname = get_fname(\n", + " fname_format = fname_format.replace(\"{blchunk}\", \"{blchunk:003}\"), kind=kind\n", " )\n", - "\n", - " out_files[kind] = create_file(fname=fname)\n" + " cross_uvd_template = create_empty_uvd(antpairs=cross_stacks[0].antpairs)\n", + " out_files[(kind, 'cross')] = create_file(fname=fname, uvd_template=cross_uvd_template)" ] }, { @@ -1313,43 +1394,43 @@ }, "outputs": [], "source": [ - "def write_data(rdc: dict, stack: lstbin.LSTStack, nbls_so_far: int, lstidx: int):\n", + "def write_data(rdc: dict, stack: lstbin.LSTStack, lstidx: int, pairs: str, template):\n", " chunk_size = stack.Nbls\n", "\n", " write = partial(\n", - " uvd_template.write_uvh5_part,\n", - " blt_inds=np.arange(nbls_so_far, nbls_so_far + chunk_size) * stackconf.n_lsts + lstidx,\n", + " template.write_uvh5_part,\n", + " blt_inds=np.arange(chunk_size) * stackconf.n_lsts + lstidx,\n", " flag_array=rdc['flags'],\n", " )\n", " \n", " write(\n", - " filename=out_files[\"LST\"],\n", + " filename=out_files[(\"LST\", pairs)],\n", " data_array=rdc[\"data\"],\n", " nsample_array=rdc[\"nsamples\"],\n", " )\n", - " print(f\"Wrote {out_files['LST']}\")\n", + " print(f\"Wrote {out_files[('LST', pairs)]}\")\n", " \n", " write(\n", - " filename=out_files[\"STD\"],\n", + " filename=out_files[(\"STD\", pairs)],\n", " data_array=rdc[\"std\"],\n", " nsample_array=rdc[\"days_binned\"],\n", " )\n", - " print(f\"Wrote {out_files['STD']}\")\n", + " print(f\"Wrote {out_files[('STD', pairs)]}\")\n", " \n", " if write_med_mad:\n", " write(\n", - " filename=out_files[\"MED\"],\n", + " filename=out_files[(\"MED\", pairs)],\n", " data_array=rdc[\"median\"],\n", " nsample_array=rdc[\"nsamples\"],\n", " )\n", - " print(f\"Wrote {out_files['MED']}\")\n", + " print(f\"Wrote {out_files[('MED', pairs)]}\")\n", " \n", " write(\n", - " filename=out_files[\"MAD\"],\n", + " filename=out_files[(\"MAD\", pairs)],\n", " data_array=rdc[\"mad\"],\n", " nsample_array=rdc[\"days_binned\"],\n", " )\n", - " print(f\"Wrote {out_files['MAD']}\")" + " print(f\"Wrote {out_files[('MAD', pairs)]}\")" ] }, { @@ -1363,8 +1444,9 @@ "source": [ "if save_lstbin_data:\n", " for lstidx in range(stackconf.n_lsts):\n", - " write_data(autos_lstavg[lstidx], auto_stacks[lstidx], 0, lstidx)\n", - " write_data(cross_lstavg[lstidx], cross_stacks[lstidx], auto_stacks[lstidx].Nbls, lstidx)\n", + " if blchunk==0:\n", + " write_data(autos_lstavg[lstidx], auto_stacks[lstidx], lstidx, 'autos', template=auto_uvd_template)\n", + " write_data(cross_lstavg[lstidx], cross_stacks[lstidx], lstidx, 'cross', template=cross_uvd_template)\n", " " ] }, @@ -2661,7 +2743,9 @@ "cell_type": "code", "execution_count": null, "id": "2002e09a-a26a-4aaa-a29b-e5d64279b0da", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "print_metadata()"