More ops/shuffeling

PSPoperations/data_handling.py

+import os
+import typing as tp
+import numpy as np
+
+# CDF library (is needed to interface with the measurement data files)
+# see https://cdf.gsfc.nasa.gov/
+os.environ["CDF_LIB"] = "/usr/local/cdf/lib"
+
+
+def cdf_slice(cdf_file, key: str):
+    """
+    DOC
+    
+    :param cdf_file:
+    :param key:
+    :return:
+    """
+    return cdf_file[key][...]
+
+
+def data_generation(cdf_file) -> tp.Dict:
+    """
+    DOC
+    
+    :param cdf_file:
+    :return:
+    """
+    data = {
+        "dqf": cdf_slice(cdf_file, key="general_flag"),
+        "epoch": cdf_slice(cdf_file, key="Epoch"),
+        "pos": cdf_slice(cdf_file, key="sc_pos_HCI"),
+        "vr": cdf_slice(cdf_file, key="vp_moment_RTN")[:, 0],
+        "dvrhi": cdf_slice(cdf_file, key="vp_moment_RTN_deltahigh")[:, 0],
+        "dvrlo": cdf_slice(cdf_file, key="vp_moment_RTN_deltalow")[:, 0],
+        "wp": cdf_slice(cdf_file, key="wp_moment"),
+        "dwphi": cdf_slice(cdf_file, key="wp_moment_deltahigh"),
+        "dwplo": cdf_slice(cdf_file, key="wp_moment_deltalow")
+    }
+    
+    return data
+
+
+def array_reduction(data_array: np.ndarray,
+                    index_list: np.ndarray) -> np.ndarray:
+    """
+    Reduction of an array by a list of indices.
+
+    :param data_array: NDARRAY,
+        Array that is to be reduced
+    :param index_list: NDARRAY,
+        List of indices that should be deleted
+    :return: NDARRAY,
+        The reduced array
+    """
+    # Choose axis=0 to delete rows and not flatten the array
+    reduced_array = np.delete(data_array, index_list, axis=0)
+
+    return reduced_array
+
+
+def full_reduction(data_dict: tp.Dict,
+                   bad_indices: np.array) -> tp.Dict:
+    """
+    DOC
+    
+    :param data_dict:
+    :param bad_indices:
+    :return:
+    """
+    for key in data_dict.keys():
+        data_dict[key] = array_reduction(data_dict[key], bad_indices)
+    
+    return data_dict

PSPoperations/data_quality.py

 import numpy as np
+import typing as tp
 
 
 def general_flag(general_flag_array: np.ndarray) -> np.ndarray:
@@ -22,19 +23,34 @@ def general_flag(general_flag_array: np.ndarray) -> np.ndarray:
     return index_list
 
 
-def array_reduction(data_array: np.ndarray,
-                    index_list: np.ndarray) -> np.ndarray:
+def meas_failed(meas_array: np.ndarray) -> np.ndarray:
     """
-    Reduction of an array by a list of indices.
+    Takes a 1D array and returns indices of failed measurements (as
+    indicated by a value of -1e30, SWEAP documentation)
     
-    :param data_array: NDARRAY,
-        Array that is to be reduced
-    :param index_list: NDARRAY,
-        List of indices that should be deleted
+    :param meas_array: NDARRAY,
+        Measurement array
     :return: NDARRAY,
-        The reduced array
+        Index array
     """
-    # Choose axis=0 to delete rows and not flatten the array
-    reduced_array = np.delete(data_array, index_list, axis=0)
+    index_array = np.where(meas_array <= -0.5e30)[0]
     
-    return reduced_array
+    return index_array
+
+
+def full_meas_eval(data_dict: tp.Dict):
+    """DOC"""
+    col_ind = []
+    
+    for key in data_dict.keys():
+        
+        # Skip unnecessary keys (maybe make this more dynamic?)
+        if key in ["pos", "dqf", "epoch"]:
+            continue
+        
+        col_ind.append(meas_failed(data_dict[key]))
+    
+    # Make an array of unique indices
+    col_ind = np.unique(np.concatenate(col_ind))
+    
+    return col_ind

PSPoperations/data_transformation.py

 import numpy as np
 import typing as tp
+# The warnings (for me) here don't seem to matter
+from astropy.constants import k_B, m_p
 
 
 def pos_cart_to_sph(position_matrix: np.ndarray) -> tp.Tuple:
@@ -28,30 +30,15 @@ def pos_cart_to_sph(position_matrix: np.ndarray) -> tp.Tuple:
 	return r, theta, phi
 
 
-def vel_cart_to_sph(vel_matrix, r, theta, phi):
+def wp_to_temp(thermal_speed: np.ndarray) -> np.ndarray:
 	"""
-	???
+	Calculate temperature from thermal speed of protons.
 	
-	:param vel_matrix:
-	:param r:
-	:param theta:
-	:param phi:
-	:return:
+	:param thermal_speed: NDARRAY,
+		Thermal speed (wp) of protons
+	:return: NDARRAY,
+		Temperature from wp = sqrt(2k_BT / m)
 	"""
-	# Sanity checks
-	assert vel_matrix.shape[1] == 3, \
-		"PLEASE CHECK POSITION MATRIX FOR COORDINATE CONVERSION!"
-	assert r.shape[0] == theta.shape[0] == phi.shape[0] == vel_matrix.shape[0], \
-		"PLEASE CHECK DIMENSIONS OF R, THETA, PHI AND VEL_MATRIX!"
-	
-	vx = vel_matrix[:, 0]
-	vy = vel_matrix[:, 1]
-	vz = vel_matrix[:, 2]
-	
-	vr = np.sin(theta) * np.cos(phi) * vx + \
-	     np.sin(theta) * np.sin(phi) * vy + \
-	     np.cos(theta) * vz
-	
-	return vr
-	
+	wp_si = thermal_speed * 1e3
 	
+	return wp_si ** 2 * m_p.value / (2 * k_B.value)

PSPoperations/plot_settings.py

+import matplotlib as mpl
+import matplotlib.pyplot as plt
+
+
+def rc_setup():
+	"""Generalized plot attributes"""
+	mpl.rcParams["xtick.direction"] = "in"
+	mpl.rcParams["ytick.direction"] = "in"
+	mpl.rcParams["xtick.minor.visible"] = "True"
+	mpl.rcParams["ytick.minor.visible"] = "True"
+
+
+def plot_r_vr(r, vr, v_ulim, v_lolim, save_ind: str = "no"):
+	"""
+	Plots radial velocity wrt heliocentric distance
+	
+	:param r: NDARRAY,
+		Distance (must be in km)
+	:param vr: NDARRAY,
+		Radial velocity (must be in km/s)
+	:param v_ulim, v_lolim: NDARRAY,
+		Upper (lower) uncertainty of vr
+	:param save_ind: STR,
+		"yes" if plot should be saved (default: no)
+	"""
+	fig, ax = plt.subplots(figsize=(10, 6))
+	plt.plot(r / 6.957e5, vr)
+	plt.fill_between(r / 6.957e5, vr + v_ulim, vr - v_lolim, color="grey",
+	                 alpha=0.25)
+	ax.set(
+		xlabel="r [R$_\\odot$]", ylabel="v$_r$ [km s$^{-1}$]",
+		ylim=(0, 800)
+	)
+	ax.grid(True, alpha=0.5)
+	
+	if save_ind.lower() == "yes":
+		plt.savefig("plot_vr.png")
+
+
+def plot_r_temp(r, T, T_ulim, T_lolim, save_ind: str = "no"):
+	"""
+	Plots logarithmic temperature wrt heliocentric distance
+	
+	:param r: NDARRAY,
+		Heliocentric distance (mst be in km)
+	:param T: NDARRAY,
+		logarithmic Temperature (Kelvin)
+	:param T_ulim, T_lolim: NDARRAY,
+		Upper (lower) uncertainty of logT
+	:param save_ind: STR,
+		"yes" if plot should be saved (default: no)
+	"""
+	fig, ax = plt.subplots(figsize=(10, 6))
+	plt.plot(r / 6.957e5, T)
+	plt.fill_between(r / 6.957e5, T_ulim, T - T_lolim,
+	                 color="grey", alpha=0.25)
+	ax.set(
+		xlabel="r [R$_\\odot$]", ylabel="T [K]", yscale="log"
+	)
+	ax.grid(True, alpha=0.5)
+	
+	if save_ind.lower() == "yes":
+		plt.savefig("plot_logT.png")

TESTING/testing.py

+import os
+import sys
+from unipath import Path
+import numpy as np
+import matplotlib.pyplot as plt
+from spacepy import pycdf
+
+# Append custom Python modules from parent directory
+sys.path.append(Path(sys.path[0]).parent)
+from PSPoperations import data_quality as dq
+from PSPoperations import data_transformation as dt
+from PSPoperations import plot_settings as ps
+from PSPoperations import data_handling as dh
+
+
+# Plot set up
+ps.rc_setup()
+
+# CDF library (is needed to interface with the measurement data files)
+# see https://cdf.gsfc.nasa.gov/
+os.environ["CDF_LIB"] = "/usr/local/cdf/lib"
+
+# open CDF file and generate faulty index array
+cdf_data = pycdf.CDF("psp_swp_spc_l3i_20211117_v02.cdf")
+data = dh.data_generation(cdf_data)
+
+# Indices of non-usable data from general flag + reduction
+bad_ind = dq.general_flag(data["dqf"])
+data = dh.full_reduction(data, bad_ind)
+
+# Additional reduction from failed measurement indices
+mf_ind = dq.full_meas_eval(data)
+data = dh.full_reduction(data, mf_ind)
+
+# Add to data dictionary (spherical coordinates)
+data["r"], data["theta"], data["phi"] = dt.pos_cart_to_sph(data["pos"])
+
+# Add to data dictionary (temperature)
+data["T"] = dt.wp_to_temp(data["wp"])
+data["dThi"] = dt.wp_to_temp(data["dwphi"])
+data["dTlo"] = dt.wp_to_temp(data["dwplo"])
+
+# vr plot
+ps.plot_r_vr(data["r"], data["vr"], data["dvrhi"], data["dvrlo"],
+             save_ind="yes")
+ps.plot_r_temp(data["r"], data["T"], data["dThi"], data["dTlo"],
+               save_ind="yes")
+
+plt.show()

main.py

+import os
+import sys
+import matplotlib.pyplot as plt
+import numpy as np
+from spacepy import pycdf
+from PSPoperations import data_quality as dq
+from PSPoperations import data_transformation as dt
+from PSPoperations import plot_settings as ps
+
+# Plot parameters
+ps.rc_setup()
+
+# CDF library (is needed to interface with the measurement data files)
+# see https://cdf.gsfc.nasa.gov/
+os.environ["CDF_LIB"] = "/usr/local/cdf/lib"
+
+
 """
 PSEUDOCODE
 
@@ -18,8 +35,71 @@ BACKGROUND
 
 OPERATIONS
 - Sort through "general flag" (only use where set to 0)
+	-- Also sort through each array to find entries with -1e30, which
+	   marks failed measurements
 - Calculate spherical heliocentric coordinates (HIC)
 - Transform to usable data parameters (vr, log_T, log_rho)
+- Generate log-file with important information
+	-- Total number of data points
+	-- Reduced number of data points
+	-- Epoch from start to finish
 """
+ENCOUNTER_NUM = "encounter_5"
+DATA_LOCATION = sys.path[0]+"/DATA/"+ENCOUNTER_NUM
+
+# SANITY CHECK IF DATA LOCATION EXISTS
+if not os.path.isdir(DATA_LOCATION):
+	print(f"\n{DATA_LOCATION} IS NOT A VALID DIRECTORY!\n")
+	sys.exit(0)
+
+
+def main():
+	# Total array setup
+	r = theta = phi = np.array([])
+	vr = dvhi = dvlo = np.array([])
+	
+	for file in os.listdir(DATA_LOCATION):
+		cdf_data = pycdf.CDF(DATA_LOCATION+"/"+file)
+		
+		dqf = cdf_data["general_flag"][...]
+		bad_ind = dq.general_flag(dqf)
+		
+		# generate reduced epoch array
+		time_file = dq.array_reduction(cdf_data["Epoch"][...], bad_ind)
+		
+		# Position
+		pos_file = dq.array_reduction(cdf_data["sc_pos_HCI"][...], bad_ind)
+		r_file, theta_file, phi_file = dt.pos_cart_to_sph(pos_file)
+		
+		# Velocity
+		vr_file = dq.array_reduction(
+			cdf_data["vp_moment_RTN"][...][:, 0],
+			bad_ind)
+		dvhi_file = dq.array_reduction(
+			cdf_data["vp_moment_RTN_deltalow"][...][:, 0],
+			bad_ind)
+		dvlo_file = dq.array_reduction(
+			cdf_data["vp_moment_RTN_deltahigh"][...][:, 0],
+			bad_ind)
+		
+		# Update total arrays
+		r = np.append(r, r_file)
+		vr = np.append(vr, vr_file)
+	
+	plt.figure()
+	plt.plot(r, vr)
+	plt.ylim(0, 1000)
+	
+	plt.show()
+
+
+def plot_setup():
+	fig_vr, ax_vr = plt.subplots()
+	ax_vr.set(
+		xlabel="r [R$_\\odot$]", ylabel="v$_r$ [km/s$^{-1}$]",
+		ylim=(0, 800)
+	)
+
 
-DATA_LOCATION = 
+if __name__ == "__main__":
+	main()