diff --git a/overlappogram/cli.py b/overlappogram/cli.py index 404fb79..b595753 100644 --- a/overlappogram/cli.py +++ b/overlappogram/cli.py @@ -49,17 +49,17 @@ def unfold(config): for alpha in config["model"]["alphas"]: for rho in config["model"]["rhos"]: - print(80*"-") + print(80 * "-") print(f"Beginning inversion for alpha={alpha}, rho={rho}.") start = time.time() - em_cube, prediction, scores, unconverged_rows = inversion.invert( - overlappogram, - config["model"], - alpha, - rho, - num_threads=config["execution"]["num_threads"], - mode_switch_thread_count=config["execution"]["mode_switch_thread_count"], - mode=MODE_MAPPING.get(config['execution']['mode'], 'invalid') + em_cube, prediction, scores, unconverged_rows, n_iter = inversion.invert( + overlappogram, + config["model"], + alpha, + rho, + num_threads=config["execution"]["num_threads"], + mode_switch_thread_count=config["execution"]["mode_switch_thread_count"], + mode=MODE_MAPPING.get(config['execution']['mode'], 'invalid') ) end = time.time() print( @@ -68,8 +68,11 @@ def unfold(config): f"seconds; {len(unconverged_rows)} unconverged rows", ) + print(f"Unconverged rows: {unconverged_rows}") + postfix = ( - "x" + str(config["inversion"]["solution_fov_width"]) + "_" + str(rho * 10) + "_" + str(alpha) + "_wpsf" + "x" + str(config["inversion"]["solution_fov_width"]) + "_" + str(rho * 10) + "_" + str( + alpha) + "_wpsf" ) save_em_cube( em_cube, @@ -90,6 +93,11 @@ def unfold(config): with open(scores_path, 'w') as f: f.write("\n".join(scores.flatten().astype(str).tolist())) + niter_path = os.path.join(config["output"]["directory"], + f"{config['output']['prefix']}_niter_{postfix}.txt") + with open(niter_path, 'w') as f: + f.write("\n".join(n_iter.flatten().astype(str).tolist())) + if config["output"]["make_spectral"]: spectral_images = create_spectrally_pure_images( [em_cube], config["paths"]["gnt"], config["inversion"]["response_dependency_list"] diff --git a/overlappogram/inversion.py b/overlappogram/inversion.py index e0d7e40..fa140b4 100644 --- a/overlappogram/inversion.py +++ b/overlappogram/inversion.py @@ -53,6 +53,7 @@ def __init__( self._row_scores: np.ndarray | None = None self._overlappogram_width: int | None = None self._overlappogram_height: int | None = None + self._n_iter: np.ndarray | None = None self._thread_count_lock = Lock() @@ -104,13 +105,15 @@ def _invert_image_row(self, row_index, chunk_index): data_out = model.predict(masked_response_function) em = model.coef_ score_data = model.score(masked_response_function, image_row) + n_iter = model.n_iter_ except ConvergenceWarning: self._unconverged_rows.append(row_index) em = np.zeros((self._num_slits * self._num_deps), dtype=np.float32) data_out = np.zeros(self._overlappogram_width, dtype=np.float32) score_data = -999 + n_iter = -1 - return row_index, em, data_out, score_data + return row_index, em, data_out, score_data, n_iter def _progress_indicator(self, future): """used in multithreading to track progress of inversion""" @@ -155,7 +158,7 @@ def _switch_to_row_inversion(self, model_config, alpha, rho, num_row_threads=50) def _collect_results(self, mode_switch_thread_count, model_config, alpha, rho): for future in concurrent.futures.as_completed(self.futures): - row_index, em, data_out, score_data = future.result() + row_index, em, data_out, score_data, n_iter = future.result() for slit_num in range(self._num_slits): if self._smooth_over == "dependence": slit_em = em[slit_num * self._num_deps : (slit_num + 1) * self._num_deps] @@ -164,6 +167,7 @@ def _collect_results(self, mode_switch_thread_count, model_config, alpha, rho): self._em_data[row_index, slit_num, :] = slit_em self._inversion_prediction[row_index, :] = data_out self._row_scores[row_index] = score_data + self._n_iter[row_index] = n_iter rows_remaining = self.total_row_count - self._completed_row_count @@ -173,7 +177,6 @@ def _collect_results(self, mode_switch_thread_count, model_config, alpha, rho): def _start_row_inversion(self, model_config, alpha, rho, num_threads): self.executors = [concurrent.futures.ThreadPoolExecutor(max_workers=num_threads)] - self.futures = {} self._models = [] for i, row_index in enumerate(range(self._detector_row_range[0], self._detector_row_range[1])): @@ -253,6 +256,7 @@ def _initialize_with_overlappogram(self, overlappogram): self._em_data = np.zeros((self._overlappogram_height, self._num_slits, self._num_deps), dtype=np.float32) self._inversion_prediction = np.zeros((self._overlappogram_height, self._overlappogram_width), dtype=np.float32) self._row_scores = np.zeros((self._overlappogram_height, 1), dtype=np.float32) + self._n_iter = np.zeros((self._overlappogram_height, 1), dtype=np.int32) def invert( self, @@ -301,4 +305,5 @@ def invert( NDCube(data=self._inversion_prediction, wcs=out_wcs, meta=self._response_meta), self._row_scores, self._unconverged_rows, + self._n_iter ) diff --git a/pyproject.toml b/pyproject.toml index 426ace5..173cb43 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ requires = ["setuptools", [project] name = "overlappogram" -version = "0.0.9" +version = "0.0.10" dependencies = ["numpy<2.0.0", "astropy", "scikit-learn", diff --git a/tests/test_inversion.py b/tests/test_inversion.py index eabce51..36e83e7 100644 --- a/tests/test_inversion.py +++ b/tests/test_inversion.py @@ -45,7 +45,7 @@ def test_inversion_runs(tmp_path, inversion_mode, is_weighted): detector_row_range=config["inversion"]["detector_row_range"], ) - em_cube, prediction, scores, unconverged_rows = inversion.invert( + em_cube, prediction, scores, unconverged_rows, niter = inversion.invert( overlappogram, config["model"], 3E-5, diff --git a/tests/test_spectral.py b/tests/test_spectral.py index 5da2fc7..d1fa01a 100644 --- a/tests/test_spectral.py +++ b/tests/test_spectral.py @@ -34,7 +34,7 @@ def test_create_spectrally_pure_images(tmp_path): detector_row_range=config["inversion"]["detector_row_range"], ) - em_cube, prediction, scores, unconverged_rows = inversion.invert( + em_cube, prediction, scores, unconverged_rows, _ = inversion.invert( overlappogram, config["model"], 3E-5,