lightstack 0.1.5__tar.gz → 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: lightstack
3
- Version: 0.1.5
3
+ Version: 0.2.0
4
4
  Summary: Tools for building and processing multi-filter astrophysical datacubes
5
5
  Author: Andressa Wille, Thallis Pessi
6
6
  License: MIT License
@@ -41,7 +41,7 @@ Requires-Dist: photutils
41
41
 
42
42
  # Lightstack
43
43
 
44
- Tools for building and processing multi-filter astrophysical datacubes.
44
+ Tools for building and processing multi-filter astrophysical datacubes.
45
45
 
46
46
  ## Features
47
47
 
@@ -49,8 +49,6 @@ Tools for building and processing multi-filter astrophysical datacubes.
49
49
  - Building photometric datacubes
50
50
  - PSF matching
51
51
 
52
- (Currently supports JWST and HST)
53
-
54
52
  ## Installation
55
53
 
56
54
  ```bash
@@ -61,4 +59,4 @@ pip install lightstack
61
59
 
62
60
  https://lightstack.readthedocs.io/en/latest/
63
61
 
64
- See the tutorial notebook!
62
+ See the tutorial notebook (in the "examples" folder)!
@@ -4,7 +4,7 @@
4
4
 
5
5
  # Lightstack
6
6
 
7
- Tools for building and processing multi-filter astrophysical datacubes.
7
+ Tools for building and processing multi-filter astrophysical datacubes.
8
8
 
9
9
  ## Features
10
10
 
@@ -12,8 +12,6 @@ Tools for building and processing multi-filter astrophysical datacubes.
12
12
  - Building photometric datacubes
13
13
  - PSF matching
14
14
 
15
- (Currently supports JWST and HST)
16
-
17
15
  ## Installation
18
16
 
19
17
  ```bash
@@ -24,4 +22,4 @@ pip install lightstack
24
22
 
25
23
  https://lightstack.readthedocs.io/en/latest/
26
24
 
27
- See the tutorial notebook!
25
+ See the tutorial notebook (in the "examples" folder)!
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "lightstack"
3
- version = "0.1.5"
3
+ version = "0.2.0"
4
4
  description = "Tools for building and processing multi-filter astrophysical datacubes"
5
5
  readme = "README.md"
6
6
  authors = [{name = "Andressa Wille"}, {name="Thallis Pessi"}]
@@ -1,42 +1,58 @@
1
1
  import numpy as np
2
2
  import os
3
+ import glob
3
4
 
4
5
  from astropy.io import fits
5
6
  from astropy.wcs import WCS
6
7
 
7
8
  from reproject import reproject_interp, reproject_exact
8
9
 
9
- from .utils import find_ext, infer_filter, get_pixel_scale_from_wcs
10
+ from .utils import find_ext, get_filter, get_pixel_scale_from_wcs
11
+
12
+
13
+ def load_image_fits(folder_path):
14
+
15
+ fits_files = []
16
+
17
+ for f in glob.glob(os.path.join(folder_path, "*.fits")):
18
+ try:
19
+ if fits.getval(f, "NAXIS") != 2:
20
+ continue
21
+ fits_files.append(f)
22
+ except Exception:
23
+ continue
24
+
25
+ return fits_files
10
26
 
11
27
 
12
28
  def align_reproject_fits(fits_list, ref_file, method="interp", crop=1):
13
29
  """
14
- Aligns and reprojects FITS images to a common WCS using a reference FITS file.
30
+ Align and optionally reproject FITS images to a common WCS.
31
+ If an image already matches the reference WCS, shape and pixel scale, reprojection is skipped automatically.
32
+ Area correction is only applied when input and output pixel scales differ.
15
33
 
16
34
  Parameters
17
35
  ----------
18
36
  fits_list : list of tuples
19
- List in the form [(fits_path, filter_name), ...].
37
+ [(fits_path, filter_name), ...]
20
38
 
21
39
  ref_file : str
22
- Path to the FITS file used as WCS reference.
40
+ Reference FITS file.
23
41
 
24
42
  method : str, optional
25
43
  Reprojection method:
26
- - "interp" (default): faster, interpolates values
27
- - "exact": slower --> but reproject_exact has precision issues with resolutions below ~0.05 arcsec, so the results may not be accurate.
44
+ - "interp" (default)
45
+ - "exact"
28
46
 
29
47
  crop : int, optional
30
- Number of pixels to remove from each border after reprojection
31
- (helps remove edge artifacts from interpolation).
48
+ Pixels to crop from borders after reprojection.
32
49
 
33
50
  Returns
34
51
  -------
35
52
  aligned_list : list of tuples
36
- List in the form [(aligned_fits_path, filter_name), ...].
53
+ [(aligned_fits_path, filter_name), ...]
37
54
  """
38
55
 
39
-
40
56
  # Choose reprojection method
41
57
  if method == "exact":
42
58
  reproj_func = reproject_exact
@@ -49,18 +65,23 @@ def align_reproject_fits(fits_list, ref_file, method="interp", crop=1):
49
65
  with fits.open(ref_file) as hdul_ref:
50
66
  ext_ref = find_ext(hdul_ref)
51
67
  if ext_ref is None:
52
- raise ValueError(f"No valid image data in reference file '{ref_file}'.")
68
+ raise ValueError(
69
+ f"No valid image data in reference file '{ref_file}'."
70
+ )
53
71
 
54
72
  ref_header = hdul_ref[ext_ref].header
55
73
  ref_wcs = WCS(ref_header)
56
74
  shape_out = hdul_ref[ext_ref].data.shape
57
75
 
58
- # Pixel scale of reference
59
76
  scale_out = get_pixel_scale_from_wcs(ref_wcs)
60
-
61
77
  aligned_list = []
62
78
 
63
- # Loop over all FITS
79
+ # Reference filter name
80
+ ref_filter_name = next(
81
+ filt for f, filt in fits_list if f == ref_file
82
+ )
83
+
84
+ # Loop through files
64
85
  for fpath, filt in fits_list:
65
86
 
66
87
  with fits.open(fpath) as hdul:
@@ -73,59 +94,99 @@ def align_reproject_fits(fits_list, ref_file, method="interp", crop=1):
73
94
  header = hdul[ext].header
74
95
  wcs_in = WCS(header)
75
96
 
76
- unit = header.get('BUNIT', 'unknown')
77
- print(f"Filter {filt}: unit = {unit}")
97
+ unit = header.get("BUNIT", "unknown")
98
+ print(f"\nFilter {filt}: unit = {unit}")
78
99
 
79
- # Pixel scale of input
80
100
  scale_in = get_pixel_scale_from_wcs(wcs_in)
101
+ print(f"Pixel scale in/out: {scale_in:.4f} -> {scale_out:.4f}")
81
102
 
82
- # Reproject
83
- data_aligned, footprint = reproj_func(
84
- (data, wcs_in),
85
- ref_wcs,
86
- shape_out=shape_out)
103
+ # Check whether reprojection is really needed
104
+ same_scale = np.isclose(scale_in, scale_out, rtol=1e-6)
105
+ same_shape = data.shape == shape_out
106
+ same_wcs = (
107
+ wcs_in.to_header_string()
108
+ == ref_wcs.to_header_string()
109
+ )
87
110
 
88
- # Area correction
89
- area_ratio = (scale_out / scale_in)**2
90
- data_aligned *= area_ratio
111
+ if same_scale and same_shape and same_wcs:
112
+ print("Already aligned with reference: skipping reprojection.")
91
113
 
92
- print(f"Pixel scale in/out: {scale_in:.4f} -> {scale_out:.4f}")
93
- print(f"Applied area correction: {area_ratio:.4f}")
114
+ data_aligned = data.copy()
115
+ wcs_out = wcs_in.deepcopy()
116
+ reprojected = False
117
+ area_corrected = False
94
118
 
95
- # Crop border
96
- if crop > 0:
97
- data_aligned = data_aligned[crop:-crop, crop:-crop]
119
+ else:
120
+ data_aligned, footprint = reproj_func(
121
+ (data, wcs_in),
122
+ ref_wcs,
123
+ shape_out=shape_out
124
+ )
98
125
 
99
- # Adjust WCS
100
126
  wcs_out = ref_wcs.deepcopy()
127
+ reprojected = True
128
+
129
+ # Area correction only if pixel scales differ
130
+ if not same_scale:
131
+ area_ratio = (scale_out / scale_in) ** 2
132
+ data_aligned *= area_ratio
133
+ print(f"Applied area correction: {area_ratio:.4f}")
134
+ area_corrected = True
135
+ else:
136
+ print("Pixel scales identical: no area correction applied.")
137
+ area_corrected = False
138
+
139
+ # Crop borders
140
+ if crop > 0:
141
+ data_aligned = data_aligned[crop:-crop, crop:-crop]
101
142
  wcs_out.wcs.crpix -= crop
102
- else:
103
- wcs_out = ref_wcs
104
143
 
105
- # Header
144
+ # Build output header
106
145
  header_aligned = wcs_out.to_header()
107
- header_aligned['BUNIT'] = unit
146
+ header_aligned["BUNIT"] = unit
108
147
 
109
- ref_filter = infer_filter(ref_file)
110
- header_aligned.add_history(
111
- f"Reprojected to {ref_filter} using {method} method (reproject package)")
112
- header_aligned.add_history(
113
- "Flux corrected by pixel area ratio")
148
+ if reprojected:
149
+ header_aligned.add_history(
150
+ f"Reprojected to {ref_filter_name} using "
151
+ f"{method} method (reproject package)"
152
+ )
153
+ else:
154
+ header_aligned.add_history(
155
+ f"No reprojection needed "
156
+ f"(already aligned to {ref_filter_name})"
157
+ )
158
+
159
+ if area_corrected:
160
+ header_aligned.add_history(
161
+ "Flux corrected by pixel area ratio"
162
+ )
163
+ else:
164
+ header_aligned.add_history(
165
+ "No area correction applied"
166
+ )
114
167
 
115
168
  if crop > 0:
116
- header_aligned.add_history(f"Cropped {crop} pixels from each border")
169
+ header_aligned.add_history(
170
+ f"Cropped {crop} pixels from each border"
171
+ )
117
172
 
118
- # Save
173
+ # Save output
119
174
  suffix = f"_aligned_{method}"
120
175
  out_name = os.path.splitext(fpath)[0] + f"{suffix}.fits"
121
176
 
122
- fits.PrimaryHDU(data_aligned, header=header_aligned).writeto(
123
- out_name, overwrite=True)
177
+ fits.PrimaryHDU(
178
+ data_aligned,
179
+ header=header_aligned
180
+ ).writeto(out_name, overwrite=True)
124
181
 
125
- # Flux sanity check
182
+ # Flux sanity check (only if reprojection happened and input not empty)
126
183
  sum_in = np.nansum(data)
127
184
  sum_out = np.nansum(data_aligned)
128
- print(f"Flux ratio (out/in): {sum_out/sum_in:.4f}")
185
+
186
+ if reprojected and sum_in != 0:
187
+ print(f"Flux ratio (out/in): {sum_out / sum_in:.4f}")
188
+ elif sum_in == 0:
189
+ print("Input image is empty: skipping flux ratio check.")
129
190
 
130
191
  aligned_list.append((out_name, filt))
131
192
  print(f"Saved: {out_name}")
@@ -133,31 +194,39 @@ def align_reproject_fits(fits_list, ref_file, method="interp", crop=1):
133
194
  return aligned_list
134
195
 
135
196
 
136
-
137
197
  def build_datacube(aligned_fits_files, reference_file, output_path):
138
198
  """
139
- Builds a 3D datacube from aligned 2D FITS images.
199
+ Build a 3D datacube from aligned 2D FITS images.
140
200
 
141
201
  Parameters
142
202
  ----------
143
203
  aligned_fits_files : list of tuples
144
- [(filename, filter_name), ...]
204
+ List in format [(filename, filter_name), ...].
205
+
145
206
  reference_file : str
146
- FITS file to define WCS and shape.
207
+ FITS file used to define WCS and output shape.
208
+
147
209
  output_path : str
148
- Path to save the 3D datacube.
149
-
210
+ Path to save the output datacube.
211
+
150
212
  Returns
151
213
  -------
152
214
  None
153
- Saves the cut datacube.
215
+ Saves the datacube to disk.
154
216
  """
155
-
156
- # Open FITS
217
+
218
+ # Safety check
219
+ if len(aligned_fits_files) == 0:
220
+ raise ValueError("No aligned FITS files provided.")
221
+
222
+ # Reference WCS
157
223
  with fits.open(reference_file) as hdul_ref:
158
224
  ext_ref = find_ext(hdul_ref)
159
225
  if ext_ref is None:
160
- raise ValueError(f"No 2D data in reference file {reference_file}.")
226
+ raise ValueError(
227
+ f"No valid 2D data in reference file {reference_file}."
228
+ )
229
+
161
230
  ref_header = hdul_ref[ext_ref].header
162
231
  ny, nx = hdul_ref[ext_ref].data.shape
163
232
  wcs_2d = WCS(ref_header, naxis=2)
@@ -166,51 +235,71 @@ def build_datacube(aligned_fits_files, reference_file, output_path):
166
235
  filter_names = []
167
236
  units = []
168
237
 
169
- # Construct cube
238
+ # Read aligned images
170
239
  for file, filt in aligned_fits_files:
240
+
171
241
  with fits.open(file) as hdul:
172
242
  ext = find_ext(hdul)
243
+
173
244
  if ext is None:
174
- print(f"No 2D data in {file}. Skipping.")
245
+ print(f"No valid 2D data in {file}. Skipping.")
175
246
  continue
247
+
176
248
  data = hdul[ext].data
249
+
177
250
  cube_images.append(data)
178
251
  filter_names.append(filt)
179
- units.append(hdul[ext].header.get('BUNIT', 'unknown'))
252
+ units.append(
253
+ hdul[ext].header.get("BUNIT", "unknown")
254
+ )
255
+
256
+ # Safety check after filtering invalid files
257
+ if len(cube_images) == 0:
258
+ raise ValueError("No valid 2D FITS images found.")
180
259
 
181
260
  cube = np.array(cube_images)
182
261
 
183
- # Write header
262
+ # Build 3D WCS
184
263
  wcs_3d = WCS(naxis=3)
264
+
265
+ # Spatial axes
185
266
  wcs_3d.wcs.crpix[0] = wcs_2d.wcs.crpix[0]
186
267
  wcs_3d.wcs.crpix[1] = wcs_2d.wcs.crpix[1]
268
+
187
269
  wcs_3d.wcs.crval[0] = wcs_2d.wcs.crval[0]
188
270
  wcs_3d.wcs.crval[1] = wcs_2d.wcs.crval[1]
271
+
189
272
  wcs_3d.wcs.cdelt[0] = wcs_2d.wcs.cdelt[0]
190
273
  wcs_3d.wcs.cdelt[1] = wcs_2d.wcs.cdelt[1]
274
+
191
275
  wcs_3d.wcs.ctype[0] = wcs_2d.wcs.ctype[0]
192
276
  wcs_3d.wcs.ctype[1] = wcs_2d.wcs.ctype[1]
277
+
193
278
  wcs_3d.wcs.cunit[0] = wcs_2d.wcs.cunit[0]
194
279
  wcs_3d.wcs.cunit[1] = wcs_2d.wcs.cunit[1]
195
280
 
281
+ # Preserve CD matrix if present
196
282
  if wcs_2d.wcs.has_cd():
197
- cd3 = np.zeros((3,3))
198
- cd3[0:2,0:2] = wcs_2d.wcs.cd.copy()
199
- cd3[2,2] = 1.0
283
+ cd3 = np.zeros((3, 3))
284
+ cd3[0:2, 0:2] = wcs_2d.wcs.cd.copy()
285
+ cd3[2, 2] = 1.0
200
286
  wcs_3d.wcs.cd = cd3
201
287
 
288
+ # Filter axis
202
289
  wcs_3d.wcs.crpix[2] = 1
203
290
  wcs_3d.wcs.crval[2] = 0
204
291
  wcs_3d.wcs.cdelt[2] = 1
205
- wcs_3d.wcs.ctype[2] = 'FILTER'
206
- wcs_3d.wcs.cunit[2] = ''
292
+ wcs_3d.wcs.ctype[2] = "FILTER"
293
+ wcs_3d.wcs.cunit[2] = ""
207
294
 
208
295
  header = wcs_3d.to_header()
209
- header['NAXIS'] = 3
210
- header['NAXIS1'] = nx
211
- header['NAXIS2'] = ny
212
- header['NAXIS3'] = len(filter_names)
213
-
296
+
297
+ header["NAXIS"] = 3
298
+ header["NAXIS1"] = nx
299
+ header["NAXIS2"] = ny
300
+ header["NAXIS3"] = len(filter_names)
301
+
302
+ # Copy HISTORY from first valid file
214
303
  with fits.open(aligned_fits_files[0][0]) as hdul0:
215
304
  ext0 = find_ext(hdul0)
216
305
  header_meta = hdul0[ext0].header
@@ -219,20 +308,26 @@ def build_datacube(aligned_fits_files, reference_file, output_path):
219
308
  for h in header_meta["HISTORY"]:
220
309
  header.add_history(h)
221
310
 
311
+ # Units
222
312
  if all(u == units[0] for u in units):
223
- header['BUNIT'] = units[0]
313
+ header["BUNIT"] = units[0]
224
314
  else:
225
- header['BUNIT'] = 'unknown'
315
+ header["BUNIT"] = "unknown"
226
316
 
317
+ # Filter metadata
227
318
  for i, filt in enumerate(filter_names):
228
- header[f'FILTER{i+1}'] = filt
229
-
230
- header['NFILTERS'] = len(filter_names)
231
- header['FILTERS'] = ",".join(filter_names)
319
+ header[f"FILTER{i+1}"] = filt
320
+
321
+ header["NFILTERS"] = len(filter_names)
322
+ header["FILTERS"] = ",".join(filter_names)
232
323
 
233
- # Save datacube
234
- fits.PrimaryHDU(cube, header=header).writeto(output_path, overwrite=True)
235
- print(f"Datacube saved at {output_path}.")
324
+ # Save cube
325
+ fits.PrimaryHDU(
326
+ cube,
327
+ header=header
328
+ ).writeto(output_path, overwrite=True)
329
+
330
+ print(f"Datacube saved at {output_path}")
236
331
 
237
332
 
238
333
  def cut_region_datacube(cube_fits_file, x_start, x_end, y_start, y_end, output_path):
@@ -247,7 +342,7 @@ def cut_region_datacube(cube_fits_file, x_start, x_end, y_start, y_end, output_p
247
342
  Pixel indices for the x axis.
248
343
  y_start, y_end : int
249
344
  Pixel indices for the y axis.
250
- output_filename : str
345
+ output_path : str
251
346
  Path to the output fits file.
252
347
 
253
348
  Returns
@@ -255,45 +350,43 @@ def cut_region_datacube(cube_fits_file, x_start, x_end, y_start, y_end, output_p
255
350
  None
256
351
  Saves the cut datacube.
257
352
  """
258
- # Open datacube
353
+
259
354
  with fits.open(cube_fits_file) as hdul:
260
355
  ext = find_ext(hdul)
261
356
  if ext is None:
262
357
  raise ValueError(f"No image data found in {cube_fits_file}")
263
358
 
264
359
  cube_data = hdul[ext].data
265
- cube_header = hdul[ext].header
360
+ cube_header = hdul[ext].header.copy()
266
361
 
267
- # Cut the data array
362
+ # Cut data
268
363
  cut_data = cube_data[:, y_start:y_end, x_start:x_end]
269
364
 
270
- # Get and update WCS
365
+ # Update WCS
271
366
  wcs_3d = WCS(cube_header)
272
367
  wcs_3d.wcs.crpix[0] -= x_start
273
368
  wcs_3d.wcs.crpix[1] -= y_start
274
369
 
275
- # Create new header with updated size and WCS
276
- new_header = wcs_3d.to_header()
277
- new_header['NAXIS'] = 3
370
+ # New header
371
+ new_header = cube_header.copy()
372
+ new_header.update(wcs_3d.to_header())
373
+
278
374
  new_header['NAXIS1'] = x_end - x_start
279
375
  new_header['NAXIS2'] = y_end - y_start
280
376
  new_header['NAXIS3'] = cube_data.shape[0]
281
377
 
282
- # Filter info
283
- for key in cube_header:
284
- if key.startswith('FILTER'):
285
- new_header[key] = cube_header[key]
286
-
287
- # Crop window
378
+ # Crop window
288
379
  new_header['XMINPIX'] = x_start
289
380
  new_header['XMAXPIX'] = x_end
290
381
  new_header['YMINPIX'] = y_start
291
382
  new_header['YMAXPIX'] = y_end
292
383
 
293
- # Save cut datacube
384
+ # Save
294
385
  fits.PrimaryHDU(cut_data, header=new_header).writeto(output_path, overwrite=True)
386
+
295
387
  print(f"Cut datacube saved to '{output_path}'")
296
-
388
+
389
+
297
390
  def build_valid_datacube(cube_fits_file, output_cube, threshold=0.0, frac_valid=0.01):
298
391
  """
299
392
  Remove empty filters in a datacube, saves the new datacube and returns the valid filter names.
@@ -336,33 +429,62 @@ def build_valid_datacube(cube_fits_file, output_cube, threshold=0.0, frac_valid=
336
429
  for i, img in enumerate(cube):
337
430
  valid = np.isfinite(img)
338
431
  n_total = valid.sum()
432
+
339
433
  if n_total == 0:
340
434
  continue
435
+
341
436
  n_above = np.sum(img[valid] > threshold)
437
+
342
438
  if (n_above / n_total) >= frac_valid:
343
439
  valid_indices.append(i)
344
-
440
+
441
+ # Stop if all filters are invalid
442
+ if len(valid_indices) == 0:
443
+ raise ValueError(
444
+ f"No valid filters found in {cube_fits_file}. "
445
+ "Try lowering threshold or frac_valid."
446
+ )
447
+
448
+ # Report removed filters
449
+ removed_filters = [
450
+ filters[i] for i in range(n_filters)
451
+ if i not in valid_indices
452
+ ]
453
+
454
+ if removed_filters:
455
+ print(f"Removed empty/invalid filters: {removed_filters}")
456
+
457
+ # Build filtered cube + valid filter names
345
458
  cube_filtered = cube[valid_indices]
346
459
  filters_valid = [filters[i] for i in valid_indices]
347
-
460
+
348
461
  # Write new header
349
462
  new_header = header.copy()
350
463
  new_header["NAXIS3"] = len(filters_valid)
351
- new_header['NFILTERS'] = len(filters_valid)
352
- new_header['FILTERS'] = ",".join(filters_valid)
464
+ new_header["NFILTERS"] = len(filters_valid)
465
+ new_header["FILTERS"] = ",".join(filters_valid)
353
466
 
354
467
  for i, f in enumerate(filters_valid):
355
468
  new_header[f"FILTER{i+1}"] = f
356
469
 
470
+ # Remove obsolete FILTER keywords
357
471
  for i in range(len(filters_valid), n_filters):
358
472
  key = f"FILTER{i+1}"
359
473
  if key in new_header:
360
474
  del new_header[key]
361
475
 
362
- # Save filtered datacube
476
+ # Add history entry
477
+ new_header.add_history(
478
+ f"Removed invalid filters "
479
+ f"(threshold={threshold}, frac_valid={frac_valid})"
480
+ )
481
+
482
+ # Save filtered datacube
363
483
  hdu = fits.PrimaryHDU(cube_filtered, header=new_header)
364
484
  hdu.writeto(output_cube, overwrite=True)
485
+
365
486
  print(f"Filtered datacube saved at '{output_cube}'")
487
+ print(f"Valid filters: {filters_valid}")
366
488
 
367
489
  return cube_filtered, filters_valid
368
490
 
@@ -378,7 +500,7 @@ def remove_filter(cube_fits_file, output_cube, filter_to_remove):
378
500
  output_cube : str
379
501
  Path to save the new datacube.
380
502
  filter_to_remove : str
381
- Name of the filter to be removed (e.g., ‘F115W’).
503
+ Name of the filter to be removed.
382
504
  """
383
505
 
384
506
  # Open FITS
@@ -393,13 +515,20 @@ def remove_filter(cube_fits_file, output_cube, filter_to_remove):
393
515
  n_filters = header.get("NAXIS3", cube.shape[0])
394
516
  filters = [header[f"FILTER{i+1}"].strip() for i in range(n_filters)]
395
517
 
396
- # Keep filters
397
- keep_indices = [i for i, f in enumerate(filters) if f != filter_to_remove]
518
+ # Keep filters except selected one
519
+ keep_indices = [i for i, f in enumerate(filters)
520
+ if f != filter_to_remove]
398
521
 
399
522
  if len(keep_indices) == len(filters):
400
- print(f"This filter '{filter_to_remove}' is not in the datacube.")
401
- else:
402
- print(f"Removing filter: {filter_to_remove}")
523
+ print(f"Filter '{filter_to_remove}' not found in datacube.")
524
+ return cube, filters
525
+
526
+ if len(keep_indices) == 0:
527
+ raise ValueError(
528
+ "Cannot remove all filters. Datacube would be empty."
529
+ )
530
+
531
+ print(f"Removing filter: {filter_to_remove}")
403
532
 
404
533
  # New cube
405
534
  cube_new = cube[keep_indices]
@@ -408,23 +537,33 @@ def remove_filter(cube_fits_file, output_cube, filter_to_remove):
408
537
  # Write new header
409
538
  new_header = header.copy()
410
539
  new_header["NAXIS3"] = len(filters_new)
411
- new_header['NFILTERS'] = len(filters_new)
412
- new_header['FILTERS'] = ",".join(filters_new)
540
+ new_header["NFILTERS"] = len(filters_new)
541
+ new_header["FILTERS"] = ",".join(filters_new)
413
542
 
414
543
  for i, f in enumerate(filters_new):
415
544
  new_header[f"FILTER{i+1}"] = f
416
545
 
546
+ # Remove old FILTER keywords
417
547
  for i in range(len(filters_new), n_filters):
418
548
  key = f"FILTER{i+1}"
419
549
  if key in new_header:
420
550
  del new_header[key]
421
551
 
552
+ # Add history
553
+ new_header.add_history(
554
+ f"Removed filter {filter_to_remove}"
555
+ )
556
+
422
557
  # Save
423
- hdu = fits.PrimaryHDU(cube_new, header=new_header)
424
- hdu.writeto(output_cube, overwrite=True)
425
- print(f"New datacube without '{filter_to_remove}' saved at {output_cube}")
426
-
558
+ fits.PrimaryHDU(cube_new, header=new_header).writeto(
559
+ output_cube,
560
+ overwrite=True
561
+ )
562
+
563
+ print(f"New datacube saved at {output_cube}")
564
+
427
565
  return cube_new, filters_new
566
+
428
567
 
429
568
  def update_cube_header(
430
569
  cube_fits_file,
@@ -514,3 +653,4 @@ def update_cube_header(
514
653
  hdul.writeto(output_file, overwrite=overwrite)
515
654
 
516
655
  print(f"Updated header saved to: {output_file}")
656
+