Comparison Plots

MODULES/Plotting/comp_plotset.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def plot_setup(indicator: str):
+	"""DOC"""
+	valid_ind = ["vr", "np", "T"]
+	assert indicator in valid_ind, f"{indicator} NOT RECOGNIZED!"
+	
+	# General values
+	fig, ax = plt.subplots(figsize=(10, 7))
+	ax.set(xlabel="Distance [R$_\\odot$]")
+	
+	if indicator == "vr":
+		ax.set(
+			ylabel="Radial velocity [km s$^{-1}$]"
+		)
+	
+	elif indicator == "np":
+		ax.set(
+			ylabel="Number density [cm$^{-3}$]",
+			yscale="log"
+		)
+	
+	elif indicator == "T":
+		ax.set(
+			ylabel="Temperature [K]",
+			yscale="log"
+		)
+	return fig, ax
+
+
+def comparison_plot(indicator: str, obs_data, sim_data, save_dir):
+	"""DOC"""
+	# Set-up plots
+	fig, ax = plot_setup(indicator)
+	
+	# GENERALIZED PLOTTING ROUTINE
+	# This is a nested class!
+	y_data = getattr(getattr(obs_data, indicator), "mean")
+	y_area = getattr(getattr(obs_data, indicator), "stddev")
+	
+	# Observational data (central)
+	ax.plot(obs_data.dist,
+	        y_data,
+	        lw=1.5,
+	        zorder=4)
+	
+	# Observational data (shaded area)
+	ax.fill_between(obs_data.dist,
+	                y1=y_data - y_area,
+	                y2=y_data + y_area,
+	                alpha=0.5,
+	                zorder=3)
+	
+	# Simulated data
+	ax.plot(sim_data.dist,
+	        getattr(sim_data, indicator),
+	        lw=2.5,
+	        zorder=5)
+	
+	plt.savefig(f"{save_dir}/{indicator}_comparison.jpg", dpi=300)
+	plt.close()
+	
\ No newline at end of file

MODULES/Plotting/general_plotset.py

+import matplotlib as mpl
+
+
+def rc_setup():
+	"""Generalized plot attributes"""
+	mpl.rcParams["xtick.direction"] = "in"
+	mpl.rcParams["xtick.labelsize"] = "large"
+	mpl.rcParams["xtick.major.width"] = 1.5
+	mpl.rcParams["xtick.minor.width"] = 1.5
+	mpl.rcParams["xtick.minor.visible"] = "True"
+	mpl.rcParams["xtick.top"] = "True"
+	
+	mpl.rcParams["ytick.direction"] = "in"
+	mpl.rcParams["ytick.labelsize"] = "large"
+	mpl.rcParams["ytick.major.width"] = 1.5
+	mpl.rcParams["ytick.minor.width"] = 1.5
+	mpl.rcParams["ytick.minor.visible"] = "True"
+	mpl.rcParams["ytick.right"] = "True"
+	
+	mpl.rcParams["axes.grid"] = "True"
+	mpl.rcParams["axes.linewidth"] = 1.5
+	mpl.rcParams["axes.labelsize"] = "large"

MODULES/Plotting/plot_settings.py → MODULES/Plotting/obs_plotset.py

-import matplotlib as mpl
+import general_plotset as gp
 import matplotlib.pyplot as plt
 import numpy as np
 
 
-def rc_setup():
-	"""Generalized plot attributes"""
-	mpl.rcParams["xtick.direction"] = "in"
-	mpl.rcParams["xtick.labelsize"] = "large"
-	mpl.rcParams["xtick.major.width"] = 1.5
-	mpl.rcParams["xtick.minor.width"] = 1.5
-	mpl.rcParams["xtick.minor.visible"] = "True"
-	mpl.rcParams["xtick.top"] = "True"
-	
-	mpl.rcParams["ytick.direction"] = "in"
-	mpl.rcParams["ytick.labelsize"] = "large"
-	mpl.rcParams["ytick.major.width"] = 1.5
-	mpl.rcParams["ytick.minor.width"] = 1.5
-	mpl.rcParams["ytick.minor.visible"] = "True"
-	mpl.rcParams["ytick.right"] = "True"
-	
-	mpl.rcParams["axes.grid"] = "True"
-	mpl.rcParams["axes.linewidth"] = 1.5
-	mpl.rcParams["axes.labelsize"] = "large"
-	
-
 def plot_setup(indicator: str):
 	"""DOC"""
 	# Make sure that oly valid indicators are used

MODULES/Statistics/data_read.py

+import numpy as np
+import astropy.constants as c
+
+
+class PSPStatData:
+	def __init__(self, filename, sim_data):
+		raw_data = np.loadtxt(filename, skiprows=1)
+		
+		# Correct for maximum distance of simulation data
+		co_index = np.where(raw_data[:, 0] >= sim_data.dist[-1])[0][0]
+		raw_data = raw_data[:co_index]
+		
+		# Positional data
+		self.dist = raw_data[:, 0]
+		
+		# Create sub-classes which contain necessary stat data
+		self.vr = StatDataSplit(raw_data, 1)
+		self.np = StatDataSplit(raw_data, 6)
+		self.T = StatDataSplit(raw_data, 11)
+
+
+class StatDataSplit:
+	"""DOC"""
+	def __init__(self, raw_data_array, first_index):
+		self.mean = raw_data_array[:, first_index]
+		self.stddev = raw_data_array[:, first_index + 1]
+		self.median = raw_data_array[:, first_index + 2]
+		self.q1 = raw_data_array[:, first_index + 3]
+		self.q3 = raw_data_array[:, first_index + 4]
+
+
+class SimMeshData:
+	def __init__(self, filename):
+		# Get first valid index
+		index = skip_nan_simdata(filename)
+		
+		# Handle simulation mesh data (skip rows integer is 1 more than index)
+		raw_data = np.loadtxt(filename, skiprows=index + 1, delimiter=",")
+		
+		# Find out if any other indices correspond to nan-values in
+		# distance and remove them
+		more_nan = np.where(np.isnan(raw_data[:, 8]))[0]
+		raw_data = np.delete(raw_data, more_nan, axis=0)
+		
+		# Assign necessary values
+		self.dist = raw_data[:, 8] / c.R_sun.value
+		self.vr = cart_to_rad_vel(raw_data) * 1e-3
+		self.np = simrho_to_rho(raw_data)
+		self.T = 10 ** raw_data[:, 3]
+		
+
+def skip_nan_simdata(filename):
+	"""DOC"""
+	# First always skipped because of header
+	raw_data = np.loadtxt(filename, skiprows=1, delimiter=",")
+	
+	# Find first valid primary index
+	first_value = 0
+	while True:
+		if np.isnan(raw_data[first_value][0]):
+			first_value += 1
+		else:
+			return first_value
+
+
+def cart_to_rad_vel(raw_data):
+	"""DOCSTRING"""
+	vx = raw_data[:, 12]
+	vy = raw_data[:, 13]
+	vz = raw_data[:, 14]
+	theta = raw_data[:, 9]
+	phi = raw_data[:, 10]
+	
+	vr = vx * np.sin(theta) * np.cos(phi) + \
+	     vy * np.sin(theta) * np.sin(phi) + \
+	     vz * np.cos(theta)
+	
+	return vr
+
+
+def simrho_to_rho(raw_data):
+	"""DOCSTRING"""
+	lrho = raw_data[:, 6]
+	rho = 10 ** lrho / c.m_p.value * 1e-6
+	
+	return rho
+	
\ No newline at end of file

TESTING/testing.py

@@ -10,7 +10,7 @@ from astropy.constants import R_sun
 sys.path.append(Path(sys.path[0]).parent)
 from MODULES.PSPops import data_quality as dq
 from MODULES.PSPops import data_transformation as dt
-from MODULES.Plotting import plot_settings as ps
+from MODULES.Plotting import obs_plotset as ps
 from MODULES.PSPops import data_handling as dh
 from MODULES.Statistics import data_binning as db
 from MODULES.Statistics import stats as st