noaa-goes-visualization/filter_FITS.py

281 lines
13 KiB
Python

import os
from collections import defaultdict
from multiprocessing import Queue, Process
import re
import traceback
import datetime
import calendar
import tqdm
import numpy as np
from matplotlib import pyplot as plt
from astropy.io import fits
from skimage.transform import hough_circle, hough_circle_peaks
from skimage.feature import canny
from skimage.filters import gaussian
from skimage.morphology import skeletonize
import cv2 as cv
def lowpriority():
""" Set the priority of the process to lowest possible."""
import sys
try:
sys.getwindowsversion()
except AttributeError:
isWindows = False
else:
isWindows = True
if isWindows:
import win32api,win32process,win32con # pywin32
pid = win32api.GetCurrentProcessId()
phandle = win32api.OpenProcess(win32con.PROCESS_ALL_ACCESS, True, pid)
win32process.SetPriorityClass(phandle, win32process.THREAD_PRIORITY_LOWEST)
# win32process.SetPriorityClass(phandle, win32process.IDLE_PRIORITY_CLASS)
# tid = win32api.GetCurrentThreadId()
# thandle = win32api.OpenThread(win32con.PROCESS_ALL_ACCESS, True, tid)
# win32process.SetThreadPriority(thandle, win32process.THREAD_MODE_BACKGROUND_BEGIN)
else:
import os
os.nice(19)
measurement_names = ["094", "131", "171", "195", "284", "304"]
thresholds = [0.050, 0.10 , 1.00, 1.40, 1.0, 2.50]
max_center_skew = 7
expected_dims = 1280
half_dims = expected_dims // 2
valid_radii = [383, 394]
avg_radius = (valid_radii[0] + valid_radii[1]) // 2
max_radius_error = 80
ratio_above_thresh_max = 0.4
ratio_above_thresh_min = 0.07
def filter_fits(work_queue):
lowpriority()
idealcircle = np.zeros((expected_dims // 2, expected_dims // 2))
cv.circle(idealcircle, (expected_dims // 4, expected_dims // 4), avg_radius // 2, 1, -1)
idealcircle_axis = np.average(idealcircle, 0)
while True:
try:
job = work_queue.get()
if job is None:
return
measurement = job.split("dr_suvi-l2-")[1].split("_")[0][2:]
idx = measurement_names.index(measurement)
try:
data = fits.getdata(job)
except IndexError: # This can happen with a blank or corrupt HDU in the .fits file
new_name = job.split(".fits")[0] + "_e.fits"
os.rename(job, new_name)
continue
assert data.shape[0] == expected_dims
assert data.shape[1] == expected_dims
data = cv.resize(data, dsize=(expected_dims // 2, expected_dims // 2), interpolation=cv.INTER_LINEAR)
filtered_data = np.copy(data)
above_thresh_indexes = data > thresholds[idx]
filtered_data[above_thresh_indexes] = thresholds[idx]
filtered_data /= thresholds[idx]
ratio_above_thresh = np.count_nonzero(above_thresh_indexes) / data.shape[0] / data.shape[1]
if ratio_above_thresh > ratio_above_thresh_max:
print(f"Exceeded ratio_above_thresh_max, possible data corruption in file: {job}")
new_name = job.split(".fits")[0] + "_e.fits"
os.rename(job, new_name)
plt.figure(figsize=[10.24, 7.68])
plt.title(f"Exceeded ratio_above_thresh_max [{ratio_above_thresh:0.2f}]")
plt.imshow(filtered_data, cmap='jet')
plt.savefig(job.split(".fits")[0] + "_e.jpg")
plt.close('all')
continue
if ratio_above_thresh < ratio_above_thresh_min:
print(f"Below ratio_above_thresh_min, possible bad data in file: {job}")
new_name = job.split(".fits")[0] + "_e.fits"
os.rename(job, new_name)
plt.figure(figsize=[10.24, 7.68])
plt.title(f"Below ratio_above_thresh_min [{ratio_above_thresh:0.2f}]")
plt.imshow(filtered_data, cmap='jet')
plt.savefig(job.split(".fits")[0] + "_e.jpg")
plt.close('all')
continue
# Decide whether the data is valid based on solar disc placement and levels.
xavg = np.average(filtered_data, 0)
yavg = np.average(filtered_data, 1)
axis_indices = list(range(half_dims))
# Center Estimation
weighted_avg_center_x = np.average(axis_indices, 0, xavg)
weighted_avg_center_y = np.average(axis_indices, 0, yavg)
good_center = True
if (half_dims // 2 - weighted_avg_center_x > max_center_skew) or (half_dims // 2 - weighted_avg_center_y > max_center_skew):
print(f"Could not find centered solar disc in file: {job}")
good_center = False
# Radius Estimation
edge_thresh = 0.98
high_edge_x = np.argmax(np.cumsum(xavg) > (np.sum(xavg) * edge_thresh))
low_edge_x = len(xavg) - np.argmax(np.cumsum(np.flip(xavg)) > (np.sum(xavg) * edge_thresh))
high_edge_y = np.argmax(np.cumsum(yavg) > (np.sum(yavg) * edge_thresh))
low_edge_y = len(yavg) - np.argmax(np.cumsum(np.flip(yavg)) > (np.sum(yavg) * edge_thresh))
good_radius = True
calc_radius_x = (high_edge_x - low_edge_x) # We do not divide by two because our image is downsized by half in each dimension
calc_radius_y = (high_edge_y - low_edge_y)
calc_radius = (calc_radius_x + calc_radius_y) / 2.0
if abs(calc_radius - avg_radius) > max_radius_error:
print(f"Could not find correct solar radius in file: {job}")
good_radius = False
# Goodness of fit estimation vs perfect disc
gof_x = np.sum(np.abs(xavg - idealcircle_axis)) / half_dims
gof_y = np.sum(np.abs(yavg - idealcircle_axis)) / half_dims
good_fit = True
if gof_x > 0.2 or gof_y > 0.2:
print(f"Could not find valid solar disc in file: {job}")
good_fit = False
if good_center and good_radius and good_fit:
# We have validated this file, rename it appropriately
new_name = job.split(".fits")[0] + "_f.fits"
os.rename(job, new_name)
else:
new_name = job.split(".fits")[0] + "_e.fits"
os.rename(job, new_name)
# Plot Results
fig, axes = plt.subplots(2, 2, width_ratios=(0.2, 1), height_ratios=(0.2, 1), gridspec_kw={"hspace":0.0, "wspace": 0.0}, figsize = [8, 8])
axes[0][0].axis('off')
axes[0][1].plot(axis_indices, xavg)
axes[0][1].plot(axis_indices, idealcircle_axis)
# axes[0][1].fill_between(axis_indices, xavg, idealcircle_axis, hatch="//", edgecolor="red", facecolor="none")
axes[1][0].plot(yavg, axis_indices)
axes[1][0].plot(idealcircle_axis, axis_indices)
axes[1][0].invert_xaxis()
axes[1][0].invert_yaxis()
axes[1][1].imshow(filtered_data, aspect='auto', vmin=0, vmax=1)
axes[1][1].vlines([low_edge_x, high_edge_x], 0, half_dims)
axes[1][1].hlines([low_edge_y, high_edge_y], 0, half_dims)
axes[1][1].scatter(half_dims // 2, half_dims // 2, marker="o", linewidths = 1, alpha = 0.75, edgecolors="black", facecolors = "none")
axes[1][1].scatter(weighted_avg_center_x, weighted_avg_center_y, marker = "x", alpha = 0.75, facecolors='red')
# Set shared x axis between imshow and xdata plots
axes[0][1].xaxis.set_ticks_position("top")
axes[0][1].yaxis.set_ticks_position("right")
axes[0][1].sharex(axes[1][1])
axes[1][0].sharey(axes[1][1])
axes[1][1].yaxis.set_ticks_position("right")
plt.savefig(job.split(".fits")[0] + "_e.jpg")
plt.close('all')
except KeyboardInterrupt:
return
except Exception as e:
print(f"Error on file: {job} - {e}")
traceback.print_exception(e)
if __name__ == "__main__":
stored_fits_dirs = [r"..\Data\goes16\l2\data\suvi-l2-ci094\2024",
r"..\Data\goes16\l2\data\suvi-l2-ci131\2024",
r"..\Data\goes16\l2\data\suvi-l2-ci171\2024",
r"..\Data\goes16\l2\data\suvi-l2-ci195\2024",
r"..\Data\goes16\l2\data\suvi-l2-ci284\2024",
r"..\Data\goes16\l2\data\suvi-l2-ci304\2024",
r"..\Data\goes18\l2\data\suvi-l2-ci094\2024",
r"..\Data\goes18\l2\data\suvi-l2-ci131\2024",
r"..\Data\goes18\l2\data\suvi-l2-ci171\2024",
r"..\Data\goes18\l2\data\suvi-l2-ci195\2024",
r"..\Data\goes18\l2\data\suvi-l2-ci284\2024",
r"..\Data\goes18\l2\data\suvi-l2-ci304\2024",]
# stored_fits_dirs = [r"..\Data\goes18\l2\data\suvi-l2-ci284\2023\01\10"]
starttime = calendar.timegm(datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc).timetuple())
stoptime = calendar.timegm(datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc).timetuple())
reprocess_errors = True
nworkers = 16
regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\.fits$"
work_queue = Queue(maxsize = nworkers)
workers = []
for i in range(nworkers):
p = Process(target = filter_fits, args = (work_queue,), daemon=True)
p.start()
workers.append(p)
lowpriority()
files_by_timestamp = defaultdict(list)
try:
for stored_fits_dir in stored_fits_dirs:
filename_tester = re.compile(regex_filename)
found_files = 0
print(f"Searching for FITS files in: {stored_fits_dir}")
for root, dirs, files in tqdm.tqdm(os.walk(stored_fits_dir), desc="Searching"):
for f in files:
if filename_tester.match(f):
file_parts = f.split("_")
file_name_end = file_parts[-1].split(".")[0]
measure_end_time = int(datetime.datetime.strptime(file_parts[4][1:16] + " +0000", "%Y%m%dT%H%M%S %z").timestamp())
if (measure_end_time >= starttime) and (measure_end_time < stoptime):
if file_name_end == "f":
if file_parts[-2] == "f": # We have an accidentally double filtered file...
new_file_name = "_".join(file_parts[:-1]) + ".fits"
os.rename(os.path.join(root,f), os.path.join(root,new_file_name))
continue # Already filtered from a previous run
elif file_name_end == "e":
if reprocess_errors:
# Rename file to remove error designation
new_file_name = "_".join(file_parts[:-1]) + ".fits"
os.rename(os.path.join(root,f), os.path.join(root,new_file_name))
# Check if error image exists and delete if needed
img_name = new_file_name.split(".fits")[0] + "_e.jpg"
try:
os.remove(os.path.join(root, img_name))
except:
pass
# Add job to queue
files_by_timestamp[measure_end_time].append(os.path.join(root, new_file_name))
else:
continue
elif file_name_end.startswith("v1-0-"): # This is the normal case for unprocessed data
files_by_timestamp[measure_end_time].append(os.path.join(root,f))
else:
# print(f"Error - Unexpected FITS file name: {f}")
pass
sorted_times = sorted(list(files_by_timestamp.keys()))
for st in tqdm.tqdm(sorted_times, desc="Filtering files"):
for f in files_by_timestamp[st]:
work_queue.put(f)
except KeyboardInterrupt:
print("Finishing current jobs and exiting")
for _ in range(nworkers):
try:
work_queue.put(None, timeout=10.0)
except:
break
for w in workers:
w.join()