diff --git a/binarycpython/utils/plot_functions.py b/binarycpython/utils/plot_functions.py
index efe85bc0662899ed52e76d3ef45c09c8104cb888..b6186bc03e1440e048fbd76b07b0f934cf10dcf4 100644
--- a/binarycpython/utils/plot_functions.py
+++ b/binarycpython/utils/plot_functions.py
@@ -16,27 +16,46 @@ TODO: Put all the plotting functions in here
"""
+import math
+
import pandas as pd
import numpy as np
+import astropy.constants as const
-# import matplotlib.pyplot as plt
+import matplotlib.pyplot as plt
from binarycpython.utils.functions import output_lines
from binarycpython.utils.run_system_wrapper import run_system
from binarycpython.utils.custom_logging_functions import binary_c_log_code
-from david_phd_functions.plotting.plot_functions import (
- plot_orbit,
- plot_masses,
- plot_HR_diagram,
-)
-
# Define the custom_logging_strings.
# These are kept to the minimum necessary for each plotting routine.
-CUSTOM_LOGGING_STRING_MASSES = ""
+CUSTOM_LOGGING_STRING_MASSES = """
+Printf("MASS_PLOTTING %30.12e %g %g %g %g\\n",
+ //
+ stardata->model.time, // 1
+
+ // masses
+ stardata->star[0].pms_mass, //
+ stardata->star[1].pms_mass, //
-CUSTOM_LOGGING_STRING_ORBIT = ""
+ stardata->star[0].mass,
+ stardata->star[1].mass
+ );
+"""
+
+CUSTOM_LOGGING_STRING_ORBIT = """
+Printf("ORBIT_PLOTTING %30.12e %g %g %g\\n",
+ //
+ stardata->model.time, // 1
+
+ // orbit elements
+ stardata->common.orbit.period, //2
+ stardata->common.orbit.separation, //3
+ stardata->common.orbit.eccentricity //4
+ );
+"""
CUSTOM_LOGGING_STRING_HR_DIAGRAM = """
Printf("HR_PLOTTING %30.12e %d %d %g %g %g %g %g %g\\n",
@@ -59,13 +78,246 @@ Printf("HR_PLOTTING %30.12e %d %d %g %g %g %g %g %g\\n",
);
"""
+def color_by_index(row, column, colors):
+ return colors[int(row[column])]
+
+def plot_HR_diagram(df, show_stellar_types=False, show_plot=True, use_astropy_values=True):
+ """
+ Function to plot the HR diagram evolution of the system. Assumes its a binary system.
+
+ For a single star see plot_HR_diagram_single
+
+ Plot shows Log luminosity on y axis,
+ log temperature on x axis (reversed)
+
+ requires:
+ time
+ stellar_type_1
+ stellar_type_2
+ radius_1
+ radius_2
+ luminosity_1
+ luminosity_2
+
+ Plots:
+ luminosity_1 vs teff_1
+ luminosity_2 vs teff_2
+
+ TODO: add HD limit
+ TODO: add lines of constant radius
+ """
+
+ # prefactor = (1/(4 * math.pi * omega_sb))**(1.0/4)
+ # prefactor_approx = 1000 * (1130)**(1.0/4) * ((R_SUN**2)/L_SUN)**(1.0/4)
+
+ if use_astropy_values:
+ R_SUN = const.R_sun.cgs.value
+ L_SUN = const.L_sun.cgs.value
+ omega_sb = const.sigma_sb.cgs.value
+ print('Using astropy values: R_SUN= {} L_SUN = {} omega_sb = {}'.format(R_SUN, L_SUN, omega_sb))
+ else:
+ R_SUN = 6.956600000000000000000000000000e+10
+ L_SUN = 3.851500000000000274321803705319e+33
+ omega_sb = 5.670352798655924736991040813194e-05
+ print('Using binary_c values: R_SUN= {} L_SUN = {} omega_sb = {}'.format(R_SUN, L_SUN, omega_sb))
+
+ prefactor = (1/(4 * math.pi * omega_sb))**(1.0/4)
+
+ if show_stellar_types:
+ fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(20, 20))
+ fig, colors = add_stellar_types_bar(df, fig, ax_index=-1, only_colorbar=True)
+ else:
+ fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(20, 20))
+
+ df = df.assign(teff_1=prefactor * ((df['luminosity_1'] * L_SUN)/((df['radius_1'] * R_SUN) ** 2)) ** (1.0/4))
+ df = df.assign(teff_2=prefactor * ((df['luminosity_2'] * L_SUN)/((df['radius_2'] * R_SUN) ** 2)) ** (1.0/4))
+
+ # # Add colors to dataframe
+ # df = df.assign(colors_1=df.apply(color_by_index, axis=1, args=('stellar_type_1', colors)))
+ # df = df.assign(colors_2=df.apply(color_by_index, axis=1, args=('stellar_type_2', colors)))
+
+ # Star 1:
+ fig.axes[0].scatter(
+ df['teff_1'],
+ df['luminosity_1'],
+ label='Star 1 (M={})'.format(df['pms_mass_1'].values.tolist()[0]),
+ # color=df['colors_1']
+ )
+
+ # Star 2:
+ fig.axes[0].scatter(
+ df['teff_2'],
+ df['luminosity_2'],
+ label='Star 2 (M={})'.format(df['pms_mass_2'].values.tolist()[0]),
+ # color=df['colors_2']
+ )
+
+ margin_fraction_x = 0.1
+ margin_fraction_y = 0.9
+
+ # Fix axes
+ min_y, max_y = df[['luminosity_1', 'luminosity_2']].min().min(), df[['luminosity_1', 'luminosity_2']].max().max()
+ min_x, max_x = df[['teff_1', 'teff_2']].min().min(), df[['teff_1', 'teff_2']].max().max()
+ fig.axes[0].set_xlim(min_x * (1-margin_fraction_x), max_x * (1 + margin_fraction_x))
+ fig.axes[0].set_ylim(min_y * (1-margin_fraction_y), max_y * (1 + margin_fraction_y))
+ fig.axes[0].set_yscale('log')
+ fig.axes[0].set_xscale('log')
+ fig.axes[0].invert_xaxis()
+
+ # Other stuff
+ fig.axes[0].set_title('HR diagram')
+ fig.axes[0].legend(loc='best')
+ fig.axes[0].set_ylabel(r'Luminosity [$L_{star}$/$L_{\odot}$]')
+ fig.axes[0].set_xlabel(r'$T_{eff}$ (K)')
+
+ # Show or return
+ if show_plot:
+ plt.show()
+ else:
+ return fig
+
+def plot_orbit(df, show_stellar_types=False, show_plot=True):
+ """
+ Function to plot the orbital elements of the system
+
+ Plots the following quantities:
+ - Orbital period
+ - Separation
+ - eccentricity
+ """
+
+ if show_stellar_types:
+ fig, ax = plt.subplots(ncols=1, nrows=4, figsize=(20, 10))
+ fig.subplots_adjust(hspace=0)
+ else:
+ fig, ax = plt.subplots(ncols=1, nrows=3, figsize=(20, 10), sharex=True)
+ fig.subplots_adjust(hspace=0)
+
+
+ #
+ fig.axes[0].plot(df['time'], df['orbital_period'], label='Orbital period')
+
+ fig.axes[1].plot(df['time'], df['separation'], label='Separation orbit')
+
+ fig.axes[2].plot(df['time'], df['eccentricity'], label='Eccentricity orbit')
+
+
+ # Make up
+ fig.axes[0].set_title('Orbital elements evolution')
+
+ fig.axes[0].legend(loc='best')
+ fig.axes[1].legend(loc='best')
+ fig.axes[2].legend(loc='best')
+
+ # fig.axes[0].set_ylim(0, 1.1*max_total_mass)
+ fig.axes[0].set_ylabel(r'Orbital period')
+ fig.axes[1].set_ylabel(r'Separation')
+ fig.axes[2].set_ylabel(r'Eccentricity')
+ fig.axes[0].set_yscale('log')
+ fig.axes[1].set_yscale('log')
+
+ fig.axes[2].set_xlabel(r'Time (Myr)')
+
+ # Show or return
+ if show_plot:
+ plt.show()
+ else:
+ return fig
+
+def plot_masses(df, show_stellar_types=False, show_plot=True):
+ """
+ Function to plot the masses of the system.
+
+ Function requires the following keys:
+ - time
+ - pms_mass_1
+ - pms_mass_2
+ - mass_1
+ - mass_2
+
+ Plots the following quantities:
+ - Total mass
+ - Mass star 1
+ - Mass star 2
+ - Pms mass 1
+ - pms mass 2
+ - pms total mass
+ (maybe?) core and env masses
+ """
+
+ if show_stellar_types:
+ fig, ax = plt.subplots(ncols=1, nrows=2, figsize=(20, 10))
+ fig.subplots_adjust(hspace=0)
+ else:
+ fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(20, 10))
+
+ max_total_mass = df['pms_mass_1'].values.tolist()[0] + df['pms_mass_2'].values.tolist()[0]
+
+ df['total_mass_temp'] = df['mass_1'] + df['mass_2']
+ df['pms_total_mass_temp'] = max_total_mass
+
+ # total mass
+ fig.axes[0].plot(
+ df['time'],
+ df['total_mass_temp'],
+ label='Total mass'
+ )
+ fig.axes[0].axhline(
+ df['pms_total_mass_temp'].values.tolist()[0],
+ label='Initial total mass',
+ linestyle='--',
+ alpha=0.5
+ )
+
+ # Mass 1
+ fig.axes[0].plot(
+ df['time'],
+ df['mass_1'],
+ label='star 1'
+ )
+ fig.axes[0].axhline(
+ df['pms_mass_1'].values.tolist()[0],
+ color='red',
+ linestyle='--',
+ linewidth=2,
+ label='Initial mass 1',
+ alpha=0.5
+ )
+
+ # mass 2
+ fig.axes[0].plot(
+ df['time'],
+ df['mass_2'],
+ color='orange',
+ label='star 2')
+ fig.axes[0].axhline(
+ df['pms_mass_2'].values.tolist()[0],
+ color='red',
+ linestyle='--',
+ linewidth=2,
+ label='Initial mass 2',
+ alpha=0.5
+ )
+
+ # Make up
+ fig.axes[0].set_title('Stellar mass evolution')
+ fig.axes[0].legend(loc='best')
+ fig.axes[0].set_ylim(0, 1.1*max_total_mass)
+ fig.axes[0].set_ylabel(r'Mass [$M_{\odot}$]')
+ fig.axes[0].set_xlabel(r'Time (Myr)')
+
+ # Show or return
+ if show_plot:
+ plt.show()
+ else:
+ return fig
+
# Define the parse functions for the plotting routines
def dummy():
"""Placeholder"""
pass
-
def parse_function_hr_diagram(output):
"""
Parsing function for the HR plotting routine
@@ -105,6 +357,67 @@ def parse_function_hr_diagram(output):
return df
+def parse_function_orbit(output):
+ """
+ Parsing function for the orbit plotting routine
+ """
+
+ values_list = []
+
+ parameters = [
+ "time",
+ "orbital_period",
+ "separation",
+ "eccentricity",
+ ]
+
+ # Go over the output.
+ for line in output_lines(output):
+ headerline = line.split()[0]
+
+ # Check the header and act accordingly
+ if headerline == "ORBIT_PLOTTING":
+ values = line.split()[1:]
+ values_list.append(values)
+
+ df = pd.DataFrame(values_list)
+ df.columns = parameters
+
+ df = df.astype(np.float64)
+
+ return df
+
+def parse_function_masses(output):
+ """
+ Parsing function for the orbit plotting routine
+ """
+
+ values_list = []
+
+ parameters = [
+ "time",
+ "pms_mass_1",
+ "pms_mass_2",
+ "mass_1",
+ "mass_2",
+ ]
+
+ # Go over the output.
+ for line in output_lines(output):
+ headerline = line.split()[0]
+
+ # Check the header and act accordingly
+ if headerline == "MASS_PLOTTING":
+ values = line.split()[1:]
+ values_list.append(values)
+
+ df = pd.DataFrame(values_list)
+ df.columns = parameters
+
+ df = df.astype(np.float64)
+
+ return df
+
def plot_system(plot_type, **kwargs):
"""
@@ -145,12 +458,12 @@ def plot_system(plot_type, **kwargs):
"mass_evolution": {
"plot_function": plot_masses,
"custom_logging_string": CUSTOM_LOGGING_STRING_MASSES,
- "parse_function": dummy,
+ "parse_function": parse_function_masses,
},
"orbit_evolution": {
"plot_function": plot_orbit,
"custom_logging_string": CUSTOM_LOGGING_STRING_ORBIT,
- "parse_function": dummy,
+ "parse_function": parse_function_orbit,
},
"hr_diagram": {
"plot_function": plot_HR_diagram,
@@ -194,7 +507,7 @@ def plot_system(plot_type, **kwargs):
run_system_arg_dict["parse_function"] = plot_types_dict[plot_type]["parse_function"]
- # Run and get the output of the parse function
+ # Run and get the output of the parse stellar_type_dict_shortfunction
binary_c_output_df = run_system(**run_system_arg_dict)
fig = plot_types_dict[plot_type]["plot_function"](
@@ -204,12 +517,18 @@ def plot_system(plot_type, **kwargs):
if not show_plot:
return fig
+# from david_phd_functions.plotting.custom_mpl_settings import load_mpl_rc
+# load_mpl_rc()
-plot_system(
- plot_type="hr_diagram",
+fig = plot_system(
+ plot_type="mass_evolution",
M_1=10,
M_2=5,
- orbital_period=100000,
+ separation=1000000,
+ orbital_period=100000000,
max_evolution_time=15000,
show_plot=True,
)
+
+# fig.axes[0].set_xlim(0, 150)
+# plt.show()
\ No newline at end of file