From 54e70c67b73e9d3e51a11491433cfc47fc8e5a10 Mon Sep 17 00:00:00 2001 From: VChristiaens Date: Thu, 18 Jul 2024 21:21:18 +0200 Subject: [PATCH] Bug fix for 4D PCA-RDI case --- vip_hci/psfsub/pca_fullfr.py | 156 ++++++++++++++--------------------- 1 file changed, 62 insertions(+), 94 deletions(-) diff --git a/vip_hci/psfsub/pca_fullfr.py b/vip_hci/psfsub/pca_fullfr.py index 48e3545c..91695c60 100644 --- a/vip_hci/psfsub/pca_fullfr.py +++ b/vip_hci/psfsub/pca_fullfr.py @@ -508,92 +508,75 @@ def pca(*all_args: List, **all_kwargs: dict): final_residuals_cube = [] recon_cube = [] medians = [] - - # (ADI+)RDI - if algo_params.cube_ref is not None: - for ch in range(nch): + table = [] + pclist = [] + grid_case = False + + # ADI or RDI + for ch in range(nch): + add_params = { + "start_time": start_time, + "cube": algo_params.cube[ch], + "ncomp": ncomp[ch], # algo_params.ncomp[ch], + "fwhm": algo_params.fwhm[ch], + "full_output": True, + } + + # RDI + if algo_params.cube_ref is not None: if algo_params.cube_ref[ch].ndim != 3: msg = "Ref cube has wrong format for 4d input cube" raise TypeError(msg) + add_params["cube_ref"] = algo_params.cube_ref[ch] - add_params = { - "start_time": start_time, - "cube": algo_params.cube[ch], - "cube_ref": algo_params.cube_ref[ch], - "ncomp": ncomp[ch], - } - func_params = setup_parameters( - params_obj=algo_params, fkt=_adi_rdi_pca, **add_params - ) - res_pca = _adi_rdi_pca( - **func_params, - **rot_options, - ) - pcs.append(res_pca[0]) - recon.append(res_pca[1]) - residuals_cube.append(res_pca[2]) - residuals_cube_.append(res_pca[3]) - ifs_adi_frames[ch] = res_pca[-1] - # ADI - else: - table = [] - pclist = [] - for ch in range(nch): - add_params = { - "start_time": start_time, - "cube": algo_params.cube[ch], - "ncomp": ncomp[ch], # algo_params.ncomp[ch], - "fwhm": algo_params.fwhm[ch], - "full_output": True, - } - func_params = setup_parameters( - params_obj=algo_params, fkt=_adi_rdi_pca, **add_params - ) - res_pca = _adi_rdi_pca( - **func_params, - **rot_options, - ) - grid_case = False - if algo_params.batch is None: - if algo_params.source_xy is not None: - # PCA grid, computing S/Ns - if isinstance(ncomp[ch], (tuple, list)): - final_residuals_cube.append(res_pca[0]) - ifs_adi_frames[ch] = res_pca[1] - table.append(res_pca[2]) - # full-frame PCA with rotation threshold - else: - recon_cube.append(res_pca[0]) - residuals_cube.append(res_pca[1]) - residuals_cube_.append(res_pca[2]) - ifs_adi_frames[ch] = res_pca[-1] + func_params = setup_parameters( + params_obj=algo_params, fkt=_adi_rdi_pca, **add_params + ) + res_pca = _adi_rdi_pca( + **func_params, + **rot_options, + ) + + if algo_params.batch is None: + if algo_params.source_xy is not None: + # PCA grid, computing S/Ns + if isinstance(ncomp[ch], (tuple, list)): + final_residuals_cube.append(res_pca[0]) + ifs_adi_frames[ch] = res_pca[1] + table.append(res_pca[2]) + # full-frame PCA with rotation threshold else: - # PCA grid - if isinstance(ncomp[ch], (tuple, list)): - ifs_adi_frames[ch] = res_pca[0] - pclist.append(res_pca[1]) - grid_case = True - # full-frame standard PCA - else: - pcs.append(res_pca[0]) - recon.append(res_pca[1]) - residuals_cube.append(res_pca[2]) - residuals_cube_.append(res_pca[3]) - ifs_adi_frames[ch] = res_pca[-1] - # full-frame incremental PCA + recon_cube.append(res_pca[0]) + residuals_cube.append(res_pca[1]) + residuals_cube_.append(res_pca[2]) + ifs_adi_frames[ch] = res_pca[-1] else: - ifs_adi_frames[ch] = res_pca[0] - pcs.append(res_pca[2]) - medians.append(res_pca[3]) - - if grid_case: - for i in range(len(ncomp[0])): - frame = cube_collapse(ifs_adi_frames[:, i], - mode=algo_params.collapse_ifs) - final_residuals_cube.append(frame) + # PCA grid + if isinstance(ncomp[ch], (tuple, list)): + ifs_adi_frames[ch] = res_pca[0] + pclist.append(res_pca[1]) + grid_case = True + # full-frame standard PCA + else: + pcs.append(res_pca[0]) + recon.append(res_pca[1]) + residuals_cube.append(res_pca[2]) + residuals_cube_.append(res_pca[3]) + ifs_adi_frames[ch] = res_pca[-1] + # full-frame incremental PCA else: - frame = cube_collapse(ifs_adi_frames, + ifs_adi_frames[ch] = res_pca[0] + pcs.append(res_pca[2]) + medians.append(res_pca[3]) + + if grid_case: + for i in range(len(ncomp[0])): + frame = cube_collapse(ifs_adi_frames[:, i], mode=algo_params.collapse_ifs) + final_residuals_cube.append(frame) + else: + frame = cube_collapse(ifs_adi_frames, + mode=algo_params.collapse_ifs) # convert to numpy arrays when relevant if len(pcs) > 0: @@ -611,22 +594,7 @@ def pca(*all_args: List, **all_kwargs: dict): if len(medians) > 0: medians = np.array(medians) - # ADI + RDI - # elif algo_params.cube_ref is not None: - # add_params = { - # "start_time": start_time, - # "full_output": True, - # } - # func_params = setup_parameters( - # params_obj=algo_params, fkt=_adi_rdi_pca, **add_params - # ) - # res_pca = _adi_rdi_pca( - # **func_params, - # **rot_options, - # ) - # pcs, recon, residuals_cube, residuals_cube_, frame = res_pca - - # ADI+RDI OR ADI. Shape of cube: (n_adi_frames, y, x) + # 3D RDI or ADI. Shape of cube: (n_adi_frames, y, x) else: add_params = { "start_time": start_time,