diff --git a/docs/pages/evolve/interface.rst b/docs/pages/evolve/interface.rst
new file mode 100644
index 00000000..eab58e16
--- /dev/null
+++ b/docs/pages/evolve/interface.rst
@@ -0,0 +1,341 @@
+**************************
+Analysing your simulations
+**************************
+
+After running your COSMIC simulations with an ``Evolve.evolve()`` call, you will have several outputs to analyse.
+The outputs are:
+
+- ``bpp`` : A pandas DataFrame containing the binary population properties at important timesteps for each binary.
+- ``bcm`` : A pandas DataFrame containing information about binaries at user-defined timesteps
+- ``initC`` : A pandas DataFrame containing the initial conditions of each binary.
+- ``kick_info`` : A pandas DataFrame containing information about natal kicks imparted to compact objects during supernovae.
+
+These outputs can be combined into a single ``COSMICOutput`` object for easier analysis and plotting.
+
+Creating a ``COSMICOutput`` object
+==================================
+
+You can create a ``COSMICOutput`` object by passing in the outputs from your evolution. Let's start by
+samples ~100 binaries and evolving them.
+
+.. ipython:: python
+
+ from cosmic.sample.initialbinarytable import InitialBinaryTable
+ from cosmic.evolve import Evolve
+ from cosmic.output import COSMICOutput
+ import matplotlib.pyplot as plt
+ import numpy as np
+
+ InitialBinaries, mass_singles, mass_binaries, n_singles, n_binaries = InitialBinaryTable.sampler(
+ 'independent', [13, 14], [13, 14], binfrac_model=0.5, primary_model='kroupa01',
+ ecc_model='sana12', porb_model='sana12', qmin=-1, SF_start=13700.0, SF_duration=0.0,
+ met=0.002, size=1000)
+
+ BSEDict = {
+ 'xi': 1.0, 'bhflag': 1, 'neta': 0.5, 'windflag': 3, 'wdflag': 1, 'alpha1': 1.0,
+ 'pts1': 0.001, 'pts3': 0.02, 'pts2': 0.01, 'epsnov': 0.001, 'hewind': 0.5, 'ck': 1000,
+ 'bwind': 0.0, 'lambdaf': 0.0, 'mxns': 3.0, 'beta': -1.0, 'tflag': 1, 'acc2': 1.5,
+ 'grflag' : 1, 'remnantflag': 4, 'ceflag': 0, 'eddfac': 1.0, 'ifflag': 0, 'bconst': 3000,
+ 'sigma': 265.0, 'gamma': -2.0, 'pisn': 45.0,
+ 'natal_kick_array' : [[-100.0,-100.0,-100.0,-100.0,0.0], [-100.0,-100.0,-100.0,-100.0,0.0]],
+ 'bhsigmafrac' : 1.0, 'polar_kick_angle' : 90,
+ 'qcrit_array' : [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],
+ 'cekickflag' : 2, 'cehestarflag' : 0, 'cemergeflag' : 0, 'ecsn' : 2.25, 'ecsn_mlow' : 1.6,
+ 'aic' : 1, 'ussn' : 0, 'sigmadiv' :-20.0, 'qcflag' : 1, 'eddlimflag' : 0,
+ 'fprimc_array' : [2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0,2.0/21.0],
+ 'bhspinflag' : 0, 'bhspinmag' : 0.0, 'rejuv_fac' : 1.0, 'rejuvflag' : 0, 'htpmb' : 1,
+ 'ST_cr' : 1, 'ST_tide' : 1, 'bdecayfac' : 1, 'rembar_massloss' : 0.5, 'kickflag' : 5,
+ 'zsun' : 0.014, 'bhms_coll_flag' : 0, 'don_lim' : -1, 'acc_lim' : -1, 'rtmsflag' : 0,
+ 'wd_mass_lim': 1
+ }
+
+ bpp, bcm, initC, kick_info = Evolve.evolve(initialbinarytable=InitialBinaries, BSEDict=BSEDict)
+
+Now we can create a ``COSMICOutput`` object quite easily (with an optional label):
+
+.. ipython:: python
+
+ output = COSMICOutput(bpp=bpp, bcm=bcm, initC=initC, kick_info=kick_info,
+ label="My First COSMIC Output")
+ print(output)
+
+This object now contains all the relevant data from our simulation in one place, and we can use its methods to analyse and plot the results.
+
+
+Saving and loading from a file
+==============================
+
+You can save a ``COSMICOutput`` object to a file for later use or sharing, and load it back when needed. The
+entire file gets saved to a single HDF5 file.
+
+.. ipython:: python
+
+ output.save('your_output_file.h5')
+
+ print(output.initC.head())
+
+ # Later, or in another script
+ loaded_output = COSMICOutput(file='your_output_file.h5')
+
+ print(loaded_output.initC.head())
+
+
+.. note::
+
+ This file will also save the version of COSMIC used to create the output. The load function will then
+ warn you if you are loading an output created with a different version of COSMIC than the one you are
+ currently using.
+
+
+Plotting population distributions
+=================================
+
+The ``COSMICOutput`` class includes several built-in plotting functions to visualize the results of your simulations.
+First, let's plot the initial mass distribution of the primary stars in our binaries.
+
+.. ipython:: python
+ :okwarning:
+ :okexcept:
+
+ fig, ax = output.plot_distribution(x_col="mass_1", when="initial", show=False);
+ @savefig initial_mass_distribution.png
+ plt.show()
+
+.. image:: initial_mass_distribution.png
+ :alt:
+ :width: 100%
+
+We could also compare this to the final mass distribution of the primary stars after evolution.
+
+.. ipython:: python
+ :okwarning:
+ :okexcept:
+
+ output.plot_distribution(x_col="mass_1", when="final", show=False);
+ @savefig final_mass_distribution.png
+ plt.show()
+
+.. image:: final_mass_distribution.png
+ :alt:
+ :width: 100%
+
+In addition to histograms, you can create scatter plots to visualize relationships between different parameters.
+
+.. ipython:: python
+ :okwarning:
+ :okexcept:
+
+ fig, ax = output.plot_distribution(
+ x_col="teff_1", y_col="lum_1", c_col="mass_1",
+ when="final", show=False
+ );
+ ax.set_xscale("log")
+ ax.set_yscale("log")
+ ax.invert_xaxis()
+
+ @savefig hrd.png
+ plt.show()
+
+.. image:: hrd.png
+ :alt:
+ :width: 100%
+
+And we also can colour the points by the stellar type of the primary star at the end of evolution and get a custom
+colour map for these ones.
+
+.. ipython:: python
+ :okwarning:
+ :okexcept:
+
+ fig, ax = output.plot_distribution(
+ x_col="mass_1", y_col="porb", c_col="kstar_1",
+ when="final", show=False
+ );
+ ax.set_xscale("log")
+ ax.set_yscale("log")
+
+ @savefig m1_porb_kstar.png
+ plt.show()
+
+.. image:: m1_porb_kstar.png
+ :alt:
+ :width: 100%
+
+Any column name from the ``bpp`` (or ``initC`` for initial conditions) DataFrames can be used for the x, y, and colour col names.
+
+Subselecting binaries
+=====================
+
+Often you'll want to focus on a specific subset of binaries from your simulation.
+You can easily subselect binaries simply by indexing the ``COSMICOutput`` object. You can index based on the
+specific ``bin_num`` of the binaries, or by using a boolean mask.
+
+Let's say you just want the first binary, or the binaries with ``bin_num`` 0, 2, and 4, or even every binary
+between 10 and 20, you just do:
+
+.. ipython:: python
+
+ first_binary = output[0]
+ some_binaries = output[[0, 2, 4]]
+ range_of_binaries = output[10:21]
+
+But perhaps you only care about binaries where the primary star starts with at least 5 solar masses and the
+binary merges during evolution. You can create a boolean mask and use that to index the ``COSMICOutput`` object:
+
+.. ipython:: python
+
+ mask = (output.initC['mass_1'] >= 5.0) & (output.final_bpp["sep"] == 0.0)
+ selected_binaries = output[mask]
+ print(selected_binaries)
+
+
+Re-running with more detailed output
+====================================
+
+Now let's say you care about one binary in particular, and you want to re-run the evolution of just that binary with more detailed output.
+You can do this easily by subselecting the binary you care about, and then calling the ``rerun_with_settings()``
+method on the resulting ``COSMICOutput`` object with a smaller ``dtp``.
+
+Let's find a binary that forms a black hole and re-run it with more detailed output:
+
+.. ipython:: python
+
+ bh = output[(output.final_bpp["kstar_1"] == 14)
+ | (output.final_bpp["kstar_2"] == 14)]
+ detailed_bh = bh.rerun_with_settings(new_settings={'dtp': 0.0}, inplace=False)
+ print(detailed_bh.bcm)
+
+Now clearly staring at the DataFrame isn't very helpful, so let's plot the evolution of the binary's separation over time.
+
+.. ipython:: python
+ :okwarning:
+ :okexcept:
+
+ bh_bin_num = detailed_bh.initC['bin_num'].iloc[0]
+ detailed_bh.plot_detailed_evolution(bin_num=bh_bin_num, show=False);
+ @savefig detailed_bh.png
+ plt.show()
+
+.. image:: detailed_bh.png
+ :alt:
+ :width: 100%
+
+This shows the full evolution, but we can ignore any time a while after the binary forms a BH
+by setting a maximum time for the x-axis.
+
+.. ipython:: python
+ :okwarning:
+ :okexcept:
+
+ formed_a_bh = (detailed_bh.bcm['kstar_1'] == 14) | (detailed_bh.bcm['kstar_2'] == 14)
+ t_max = detailed_bh.bcm['tphys'][formed_a_bh].min() + 10.0 # 10 Myr after first BH forms
+
+ detailed_bh.plot_detailed_evolution(bin_num=bh_bin_num,
+ t_max=t_max, show=False);
+ @savefig detailed_bh_limited.png
+ plt.show()
+
+.. image:: detailed_bh_limited.png
+ :alt:
+ :width: 100%
+
+
+Re-running with new physics settings
+====================================
+
+You can also re-run the evolution of your binaries with different physics settings, but the same initial conditions.
+This is useful for testing how different assumptions impact your results. You can do this using the ``rerun_with_settings()`` method, passing in a dictionary of
+new settings. For example, let's say we want to see how changing the common envelope efficiency affects our results.
+
+.. ipython:: python
+
+ ce_alpha_10 = output.rerun_with_settings(
+ new_settings={'alpha1': 10}, inplace=False
+ )
+ n_merger_original = len(output.final_bpp[output.final_bpp["sep"] == 0.0])
+ n_merger_ce_alpha_10 = len(ce_alpha_10.final_bpp[ce_alpha_10.final_bpp["sep"] == 0.0])
+ print(f"Original number of mergers: {n_merger_original}")
+ print(f"Number of mergers with ceflag=10: {n_merger_ce_alpha_10}")
+
+This shows how making common-envelope evolution more efficient allows more binaries to survive and avoid merging.
+
+
+Additionally, if your new settings affect natal kicks (e.g., changing remnant mass prescriptions), you can
+also choose to reset the natal kicks when re-running the evolution. This ensures that the kicks are sampled
+according to the new physics settings. You can do this by setting the ``reset_kicks`` parameter to ``True`` in the ``rerun_with_settings()`` method.
+
+.. ipython:: python
+
+ kickflag_1 = output.rerun_with_settings(
+ new_settings={'kickflag': 1}, reset_kicks=True, inplace=False
+ )
+ print(f"Average kick velocity with original settings: {output.kick_info['natal_kick'].mean():1.2f} km/s")
+ print(f"Average kick velocity with kickflag=1 and reset kicks: {kickflag_1.kick_info['natal_kick'].mean():1.2f} km/s")
+
+
+Examining output from `cosmic-pop`
+==================================
+
+If you have run a population synthesis simulation using ``cosmic-pop``, you can also load the output
+from that simulation into a ``COSMICPopOutput`` object. This object extends the functionality of ``COSMICOutput``
+to handle the binary and single star populations from ``cosmic-pop``, as well as storing information about the
+convergence of the simulation.
+
+Let's say that you've saved a ``cosmic-pop`` simulation to a file called ``dat_kstar1_13_14_kstar2_13_14_SFstart_13700.0_SFduration_0.0_metallicity_0.02.h5``.
+You can load this file into a ``COSMICPopOutput`` object like so:
+
+.. code-block:: python
+
+ from cosmic.output import COSMICPopOutput
+
+ pop_output = COSMICPopOutput(file='dat_kstar1_13_14_kstar2_13_14_SFstart_13700.0_SFduration_0.0_metallicity_0.02.h5')
+ print(pop_output)
+
+
+This file contains both binary and single star populations, which you can access via the ``output`` and ``singles_output`` attributes, respectively.
+These are both ``COSMICOutput`` objects, so you can use all the same methods and attributes on them as well (e.g., plotting distributions, subselecting binaries, re-running with new settings, etc.).
+
+.. code-block:: python
+
+ # the full bpp table for binaries
+ print(pop_output.output.bpp)
+
+If you set ``keep_singles=False`` when running ``cosmic-pop``, the ``singles_output`` attribute will be ``None``.
+
+You can access the usual cosmic-pop outputs as attributes of the class, such as:
+
+.. code-block:: python
+
+ print(pop_output.conv)
+ print(pop_output.match)
+ print(pop_output.n_binaries)
+ print(pop_output.mass_stars)
+
+Combining binary and single star outputs
+----------------------------------------
+
+If you want to combine the binary and single star outputs into a single ``COSMICOutput`` object, you can use the ``to_combined_output()`` method.
+This will concatenate the binary and single star DataFrames together. This can let you more easily analyse the full stellar population from your simulation.
+
+.. code-block:: python
+
+ combined_output = pop_output.to_combined_output()
+
+ # for example, we could select the BHBH binaries and re-run them with more detailed output
+ bhbh_binaries = combined_output[
+ (combined_output.final_bpp["kstar_1"] == 14) &
+ (combined_output.final_bpp["kstar_2"] == 14)
+ ]
+ detailed_bhbh = bhbh_binaries.rerun_with_settings(new_settings={'dtp': 0.1}, inplace=False)
+ detailed_bhbh.plot_detailed_evolution(
+ bin_num=detailed_bhbh.initC['bin_num'].iloc[0],
+ t_max=100
+ );
+
+
+Wrapping up
+===========
+
+The ``COSMICOutput`` class provides a convenient way to manage and analyse the results of your COSMIC simulations.
+With built-in methods for plotting, subselecting binaries, and re-running evolutions with different settings,
+it will hopefully make your analysis workflow smoother and get you to your interesting results faster!
\ No newline at end of file
diff --git a/docs/pages/examples.rst b/docs/pages/examples.rst
index d95b737e..c35e30ff 100644
--- a/docs/pages/examples.rst
+++ b/docs/pages/examples.rst
@@ -16,4 +16,5 @@ COSMIC can evolve binaries for several different use cases. Go through the examp
evolve/resolution
evolve/rerun
evolve/restart
+ evolve/interface
diff --git a/src/cosmic/evolve.py b/src/cosmic/evolve.py
index 23f4092a..fbbb1fcf 100644
--- a/src/cosmic/evolve.py
+++ b/src/cosmic/evolve.py
@@ -282,57 +282,54 @@ def evolve(self, initialbinarytable, pool=None, bpp_columns=None, bcm_columns=No
if 'bin_num' not in initialbinarytable.keys():
initialbinarytable = initialbinarytable.assign(bin_num=np.arange(idx, idx + len(initialbinarytable)))
+ # go through each item in the BSEDict and update the initialbinarytable
new_cols = {}
n = len(initialbinarytable)
idx = initialbinarytable.index
-
for k, v in list(BSEDict.items()):
+ # warn the user if they are overwriting a value
if k in initialbinarytable.columns:
warnings.warn(
"The value for {0} in initial binary table is being overwritten by the value of {0} "
"from either the params file or the BSEDict.".format(k)
)
+ # handle special cases where we need to expand arrays into multiple columns
if k == 'natal_kick_array':
- new_cols['natal_kick_array'] = pd.Series(
- [BSEDict['natal_kick_array']] * n, index=idx
- )
-
+ initialbinarytable["natal_kick_array"] = [BSEDict['natal_kick_array']] * n
for j, column_name in enumerate(NATAL_KICK_COLUMNS):
for sn in range(2):
col = f"{column_name}_{sn+1}"
- new_cols[col] = pd.Series(
- [BSEDict['natal_kick_array'][sn][j]] * n,
- index=idx,
- )
+ if col in initialbinarytable.columns:
+ initialbinarytable[col] = BSEDict['natal_kick_array'][sn][j]
+ else:
+ new_cols[col] = BSEDict['natal_kick_array'][sn][j]
elif k == 'qcrit_array':
- new_cols['qcrit_array'] = pd.Series(
- [BSEDict['qcrit_array']] * n, index=idx
- )
-
+ initialbinarytable["qcrit_array"] = [BSEDict['qcrit_array']] * n
for kstar in range(16):
- new_cols[f"qcrit_{kstar}"] = pd.Series(
- [BSEDict['qcrit_array'][kstar]] * n,
- index=idx,
- )
+ col = f"qcrit_{kstar}"
+ if col in initialbinarytable.columns:
+ initialbinarytable[col] = BSEDict['qcrit_array'][kstar]
+ else:
+ new_cols[col] = BSEDict['qcrit_array'][kstar]
elif k == 'fprimc_array':
- new_cols['fprimc_array'] = pd.Series(
- [BSEDict['fprimc_array']] * n, index=idx
- )
-
+ initialbinarytable["fprimc_array"] = [BSEDict['fprimc_array']] * n
for kstar in range(16):
- new_cols[f"fprimc_{kstar}"] = pd.Series(
- [BSEDict['fprimc_array'][kstar]] * n,
- index=idx,
- )
-
+ col = f"fprimc_{kstar}"
+ if col in initialbinarytable.columns:
+ initialbinarytable[col] = BSEDict['fprimc_array'][kstar]
+ else:
+ new_cols[col] = BSEDict['fprimc_array'][kstar]
else:
- # scalar or array-like; pandas will broadcast scalars automatically
- new_cols[k] = v
+ # base case: if it's present, overwrite, if not, add to a list of new columns (see below)
+ if k in initialbinarytable.columns:
+ initialbinarytable[k] = v
+ else:
+ new_cols[k] = v
- # concat once (fast, no fragmentation)
+ # for columns that are new to the initial binary table, concat once
if new_cols:
new_df = pd.DataFrame(new_cols, index=idx)
initialbinarytable = pd.concat([initialbinarytable, new_df], axis=1)
@@ -437,9 +434,15 @@ def evolve(self, initialbinarytable, pool=None, bpp_columns=None, bcm_columns=No
# update initial table with sampled kicks
to_add = {}
for idx, column in enumerate(FLATTENED_NATAL_KICK_COLUMNS):
- to_add[column] = natal_kick_arrays[:, 0, idx]
- natal_kick_df = pd.DataFrame(to_add, index=initialbinarytable.index)
- initialbinarytable = pd.concat([initialbinarytable, natal_kick_df], axis=1)
+ if column not in initialbinarytable.columns:
+ to_add[column] = natal_kick_arrays[:, 0, idx]
+ else:
+ initialbinarytable[column] = natal_kick_arrays[:, 0, idx]
+
+ # if kicks weren't already present, add them
+ if to_add:
+ natal_kick_df = pd.DataFrame(to_add, index=initialbinarytable.index)
+ initialbinarytable = pd.concat([initialbinarytable, natal_kick_df], axis=1)
kick_info = pd.DataFrame(kick_info_arrays,
columns=KICK_COLUMNS,
diff --git a/src/cosmic/meson.build b/src/cosmic/meson.build
index 0bf13fc6..899314bd 100644
--- a/src/cosmic/meson.build
+++ b/src/cosmic/meson.build
@@ -8,7 +8,8 @@ python_sources = [
'get_commit_hash.py',
'Match.py',
'plotting.py',
- 'utils.py'
+ 'utils.py',
+ 'output.py',
]
py3.install_sources(
diff --git a/src/cosmic/output.py b/src/cosmic/output.py
new file mode 100644
index 00000000..2e7bc78b
--- /dev/null
+++ b/src/cosmic/output.py
@@ -0,0 +1,466 @@
+import json
+import pandas as pd
+import h5py as h5
+from cosmic.evolve import Evolve
+from cosmic._version import __version__
+from cosmic.plotting import plot_binary_evol
+import matplotlib.pyplot as plt
+from matplotlib.colors import ListedColormap, BoundaryNorm
+import numpy as np
+import warnings
+
+
+__all__ = ['COSMICOutput', 'COSMICPopOutput', 'save_initC', 'load_initC']
+
+
+kstar_translator = [
+ {'long': 'Main Sequence (Low mass)', 'short': 'MS < 0.7', 'colour': (0.996078, 0.843476, 0.469158, 1.0)},
+ {'long': 'Main Sequence', 'short': 'MS', 'colour': (0.996078, 0.843476, 0.469158, 1.0)},
+ {'long': 'Hertzsprung Gap', 'short': 'HG', 'colour': (0.939608, 0.471373, 0.094902, 1.0)},
+ {'long': 'First Giant Branch', 'short': 'FGB', 'colour': (0.716186, 0.833203, 0.916155, 1.0)},
+ {'long': 'Core Helium Burning', 'short': 'CHeB', 'colour': (0.29098, 0.59451, 0.78902, 1.0)},
+ {'long': 'Early AGB', 'short': 'EAGB', 'colour': (0.294902, 0.690196, 0.384314, 1.0)},
+ {'long': 'Thermally Pulsing AGB', 'short': 'TPAGB',
+ 'colour': (0.723122, 0.889612, 0.697178, 1.0)},
+ {'long': 'Helium Main Sequence', 'short': 'HeMS', 'colour': (0.254627, 0.013882, 0.615419, 1.0)},
+ {'long': 'Helium Hertsprung Gap', 'short': 'HeHG', 'colour': (0.562738, 0.051545, 0.641509, 1.0)},
+ {'long': 'Helium Giant Branch', 'short': 'HeGB', 'colour': (0.798216, 0.280197, 0.469538, 1.0)},
+ {'long': 'Helium White Dwarf', 'short': 'HeWD', 'colour': (0.368166, 0.232828, 0.148275, 1.0)},
+ {'long': 'Carbon/Oxygen White Dwarf', 'short': 'COWD', 'colour': (0.620069, 0.392132, 0.249725, 1.0)},
+ {'long': 'Oxygen/Neon White Dwarf', 'short': 'ONeWD', 'colour': (0.867128, 0.548372, 0.349225, 1.0)},
+ {'long': 'Neutron Star', 'short': 'NS', 'colour': (0.501961, 0.501961, 0.501961, 1.0)},
+ {'long': 'Black Hole', 'short': 'BH', 'colour': (0.0, 0.0, 0.0, 1.0)},
+ {'long': 'Massless Remnant', 'short': 'MR', 'colour': "white"},
+ {'long': 'Chemically Homogeneous', 'short': 'CHE', 'colour': (0.647059, 0.164706, 0.164706, 1.0)}
+]
+
+
+class COSMICOutput:
+ def __init__(self, bpp=None, bcm=None, initC=None, kick_info=None, file=None, label=None,
+ file_key_suffix=''):
+ """Container for COSMIC output data components.
+
+ Can be initialized either from data components directly or by loading from an HDF5 file.
+
+ Parameters
+ ----------
+ bpp : `pandas.DataFrame`, optional
+ Important evolution timestep table, by default None
+ bcm : `pandas.DataFrame`, optional
+ User-defined timestep table, by default None
+ initC : `pandas.DataFrame`, optional
+ Initial conditions table, by default None
+ kick_info : `pandas.DataFrame`, optional
+ Natal kick information table, by default None
+ file : `str`, optional
+ Filename/path to HDF5 file to load data from, by default None
+ label : `str`, optional
+ Optional label for the output instance, by default None
+ file_key_suffix : `str`, optional
+ Suffix to append to dataset keys when loading from file, by default ''. E.g. if set to '_singles',
+ datasets 'bpp_singles', 'bcm_singles', etc. will be loaded as bpp, bcm, etc.
+
+ Raises
+ ------
+ ValueError
+ If neither file nor all data components are provided.
+ """
+ # require that either file is given or all data components are given
+ if file is None and (bpp is None or bcm is None or initC is None or kick_info is None):
+ raise ValueError("Either file or all data components (bpp, bcm, initC, kick_info) must be provided.")
+ if file is not None:
+ self.bpp = pd.read_hdf(file, key=f'bpp{file_key_suffix}')
+ self.bcm = pd.read_hdf(file, key=f'bcm{file_key_suffix}')
+ self.initC = load_initC(file, key=f'initC{file_key_suffix}',
+ settings_key=f'initC_{file_key_suffix}_settings')
+ self.kick_info = pd.read_hdf(file, key=f'kick_info{file_key_suffix}')
+ with h5.File(file, 'r') as f:
+ file_version = f.attrs.get('COSMIC_version', 'unknown')
+ label = f.attrs.get('label', '')
+ self.label = label if label != '' else None
+ if file_version != __version__:
+ warnings.warn(f"You have loaded COSMICOutput from a file that was run using COSMIC version {file_version}, "
+ f"but the current version is {__version__}. "
+ "There may be compatibility issues, or differences in output when rerunning, be sure to check the changelog.", UserWarning)
+ else:
+ self.bpp = bpp
+ self.bcm = bcm
+ self.initC = initC
+ self.kick_info = kick_info
+ self.label = label if label is not None else None
+
+ def __len__(self):
+ return len(self.initC)
+
+ def __repr__(self):
+ return f''
+
+ def __getitem__(self, key):
+ """Subselect binaries by bin_num across all data components.
+ Keys can be integers or lists/arrays of integers or slices.
+ If the key is an array of bools, mask initC to get the corresponding bin_nums."""
+ # convert key to list of bin_nums, regardless of input type
+ if isinstance(key, int):
+ key = [key]
+ elif isinstance(key, slice):
+ key = self.initC['bin_num'].iloc[key].tolist()
+ elif isinstance(key, (pd.Series, list, np.ndarray)) and len(key) == len(self.initC) and isinstance(key[0], (bool, np.bool_)):
+ if not key.any():
+ raise IndexError("Boolean mask resulted in zero selected binaries.")
+ key = self.initC['bin_num'][key].tolist()
+ # otherwise, reject invalid types
+ elif not isinstance(key, (list, np.ndarray, pd.Series)):
+ raise TypeError("Key must be an int, slice, list/array of ints, or boolean mask.")
+
+ bpp_subset = self.bpp[self.bpp['bin_num'].isin(key)]
+ bcm_subset = self.bcm[self.bcm['bin_num'].isin(key)]
+ initC_subset = self.initC[self.initC['bin_num'].isin(key)]
+ kick_info_subset = self.kick_info[self.kick_info['bin_num'].isin(key)]
+ return COSMICOutput(bpp=bpp_subset, bcm=bcm_subset, initC=initC_subset, kick_info=kick_info_subset, label=self.label)
+
+ @property
+ def final_bpp(self):
+ """Get the final timestep for each binary from the bpp table.
+
+ Returns
+ -------
+ final_bpp : `pandas.DataFrame`
+ DataFrame containing only the final timestep for each binary.
+ """
+ return self.bpp.drop_duplicates(subset='bin_num', keep='last')
+
+ def save(self, output_file):
+ """Save all data components to an HDF5 file
+
+ Parameters
+ ----------
+ output_file : `str`
+ Filename/path to the HDF5 file
+ """
+ self.bpp.to_hdf(output_file, key='bpp')
+ self.bcm.to_hdf(output_file, key='bcm')
+ save_initC(output_file, self.initC, key='initC', settings_key='initC_settings')
+ self.kick_info.to_hdf(output_file, key='kick_info')
+ with h5.File(output_file, 'a') as f:
+ f.attrs['COSMIC_version'] = __version__
+ f.attrs['label'] = self.label if self.label is not None else ''
+
+ def rerun_with_settings(self, new_settings, reset_kicks=False, inplace=False):
+ """Rerun the simulation with new settings.
+
+ Parameters
+ ----------
+ new_settings : `dict`
+ Dictionary of new settings to apply. Any setting not included will retain its original value.
+ reset_kicks : `bool`, optional
+ If True, reset natal kicks to be randomly sampled again.
+ If False, retain original kicks. By default False.
+ (You may want to reset the kicks if changing settings that affect remnant masses or
+ kick distribution.)
+ inplace : `bool`, optional
+ If True, update the current instance. If False, return a new instance. By default False.
+
+ Returns
+ -------
+ new_output : `COSMICOutput`
+ New COSMICOutput instance with updated simulation results (only if inplace is False).
+ """
+ # merge new settings with existing initC
+ updated_initC = self.initC.copy()
+ for key, value in new_settings.items():
+ if key in updated_initC.columns:
+ updated_initC[key] = value
+ else:
+ raise KeyError(f"Setting '{key}' not found in initC columns.")
+
+ # reset kicks if requested
+ if reset_kicks:
+ kick_cols = ["natal_kick_1", "natal_kick_2", "phi_1", "phi_2", "theta_1", "theta_2",
+ "mean_anomaly_1", "mean_anomaly_2"]
+ for col in kick_cols:
+ updated_initC[col] = -100.0
+ elif 'kickflag' in new_settings or 'remnantflag' in new_settings:
+ warnings.warn(
+ "You have changed 'kickflag' or 'remnantflag' without resetting kicks. "
+ "This may lead to inconsistent results if the kick distribution or remnant masses have changed. "
+ "Consider setting reset_kicks=True.", UserWarning
+ )
+
+ # re-run the simulation
+ new_bpp, new_bcm, new_initC, new_kick_info = Evolve.evolve(initialbinarytable=updated_initC)
+
+ if inplace:
+ self.bpp = new_bpp
+ self.bcm = new_bcm
+ self.initC = new_initC
+ self.kick_info = new_kick_info
+ else:
+ return COSMICOutput(bpp=new_bpp, bcm=new_bcm, initC=new_initC, kick_info=new_kick_info)
+
+
+ def plot_detailed_evolution(self, bin_num, show=True, **kwargs):
+ """Plot detailed evolution for a specific binary.
+
+ Parameters
+ ----------
+ bin_num : `int`
+ Index of the binary to plot.
+ **kwargs :
+ Additional keyword arguments passed to the plotting function (plotting.plot_binary_evol).
+ """
+ # check the bin_num is in the bcm
+ if bin_num not in self.bcm['bin_num'].values:
+ raise ValueError(f"bin_num {bin_num} not found in bcm table.")
+
+ # warn if bcm has only two entries for this binary
+ bcm_subset = self.bcm[self.bcm['bin_num'] == bin_num]
+ if len(bcm_subset) <= 2:
+ warnings.warn(
+ f"bcm table for bin_num {bin_num} has only {len(bcm_subset)} entries. Detailed evolution "
+ "plot may be uninformative. You should set dtp, or timestep_conditions, to increase the "
+ "number of timesteps in the bcm table.", UserWarning
+ )
+
+ if "ktype_kwargs" not in kwargs:
+ kwargs["ktype_kwargs"] = {'k_type_colors': [kstar_translator[k]["colour"] for k in range(len(kstar_translator))]}
+ fig = plot_binary_evol(self.bcm.loc[bin_num], **kwargs)
+ if show:
+ plt.show()
+ return fig
+
+
+ def plot_distribution(self, x_col, y_col=None, c_col=None, when='final',
+ fig=None, ax=None, show=True,
+ xlabel='auto', ylabel='auto', clabel='auto', **kwargs):
+ """Plot distribution of binaries in specified columns.
+
+ Plots can be histograms (if only x_col is given) or scatter plots (if both x_col and y_col are given).
+ Optionally, colour coding can be applied using c_col.
+
+ Parameters
+ ----------
+ x_col : `str`
+ Column name for x-axis.
+ y_col : `str`, optional
+ Column name for y-axis. If None, a histogram will be plotted. By default None.
+ c_col : `str`, optional
+ Column name for colour coding. By default None.
+ when : `str`, optional
+ When to take the values from: 'initial' or 'final'. By default 'final'.
+ fig : `matplotlib.figure.Figure`, optional
+ Figure to plot on. If None, a new figure is created. By default None.
+ ax : `matplotlib.axes.Axes`, optional
+ Axes to plot on. If None, new axes are created. By default None.
+ show : `bool`, optional
+ If True, display the plot immediately. By default True.
+ xlabel : `str`, optional
+ Label for x-axis. If 'auto', uses the column name. By default 'auto'.
+ ylabel : `str`, optional
+ Label for y-axis. If 'auto', uses the column name or 'Count' for histogram. By default 'auto'.
+ clabel : `str`, optional
+ Label for colorbar. If 'auto', uses the column name. By default 'auto
+ **kwargs :
+ Additional keyword arguments passed to the plotting function.
+
+ Returns
+ -------
+ fig : `matplotlib.figure.Figure`
+ The figure containing the plot.
+ ax : `matplotlib.axes.Axes`
+ The axes containing the plot.
+ """
+ if fig is None or ax is None:
+ fig, ax = plt.subplots()
+
+ if when == 'initial':
+ data = self.initC
+ elif when == 'final':
+ data = self.bpp.drop_duplicates(subset='bin_num', keep='last')
+ else:
+ raise ValueError("Parameter 'when' must be either 'initial' or 'final'.")
+
+ if xlabel == 'auto':
+ xlabel = x_col
+ if ylabel == 'auto':
+ ylabel = y_col if y_col is not None else 'Count'
+ if clabel == 'auto' and c_col is not None:
+ clabel = c_col
+
+ if y_col is None:
+ # histogram
+ ax.hist(data[x_col], bins=kwargs.get('bins', "fd"),
+ color=kwargs.get('color', "tab:blue"), **kwargs)
+ ax.set(
+ xlabel=xlabel,
+ ylabel=ylabel,
+ )
+ else:
+ # scatter plot
+ c = data[c_col] if c_col is not None else kwargs.get('color', None)
+ if c_col == 'kstar_1' or c_col == 'kstar_2':
+ c = data[c_col].map(lambda k: kstar_translator[k]['colour'])
+ sc = ax.scatter(data[x_col], data[y_col],
+ c=c,
+ **kwargs)
+ ax.set(
+ xlabel=xlabel,
+ ylabel=ylabel,
+ )
+ if c_col is not None and c_col not in ['kstar_1', 'kstar_2']:
+ cbar = fig.colorbar(sc, ax=ax)
+ cbar.set_label(clabel)
+ elif c_col is not None:
+ # extract colours and labels
+ colours = [entry["colour"] for entry in kstar_translator[1:-2]]
+ labels = [entry["short"] for entry in kstar_translator[1:-2]]
+
+ # create colormap
+ cmap = ListedColormap(colours)
+ bounds = np.arange(len(colours) + 1)
+ norm = BoundaryNorm(bounds, cmap.N)
+
+ cb = plt.colorbar(
+ mappable=plt.cm.ScalarMappable(norm=norm, cmap=cmap),
+ ticks=np.arange(len(colours)) + 0.5,
+ boundaries=bounds,
+ ax=ax
+ )
+ cb.ax.set_yticklabels(labels)
+ cb.set_label(clabel)
+
+ if show:
+ plt.show()
+ return fig, ax
+
+
+class COSMICPopOutput():
+ def __init__(self, file, label=None):
+ # read in convergence tables and totals
+ keys = ['conv', 'idx', 'match', 'mass_binaries', 'mass_singles',
+ 'n_binaries', 'n_singles', 'mass_stars', 'n_stars']
+ for key in keys:
+ setattr(self, key, pd.read_hdf(file, key=key))
+
+ # load config back from JSON storage
+ with h5.File(file, 'r') as f:
+ self.config = json.loads(f['config'][()])
+
+ # create a COSMICOutput for the binaries, and optionally for the singles
+ self.output = COSMICOutput(file=file, label=label + ' [binaries]' if label is not None else None)
+ singles = "keep_singles" in self.config["sampling"] and self.config["sampling"]["keep_singles"]
+ self.singles_output = COSMICOutput(
+ file=file, label=label + ' [singles]' if label is not None else None,
+ file_key_suffix='_singles'
+ ) if singles else None
+ self.label = label
+
+ def __repr__(self):
+ r = f''
+ if self.singles_output is not None:
+ r = r[:-1] + f', {len(self.singles_output)} singles>'
+ return r
+
+ def __len__(self):
+ return len(self.conv)
+
+ def to_combined_output(self):
+ """Combine binaries and singles into a single COSMICOutput instance.
+
+ Returns
+ -------
+ combined_output : `COSMICOutput`
+ COSMICOutput instance containing both binaries and singles.
+
+ Raises
+ ------
+ ValueError
+ If singles output is not available.
+ """
+ if self.singles_output is None:
+ raise ValueError("Singles output is not available in this COSMICPopOutput instance.")
+
+ bpp = pd.concat([self.output.bpp, self.singles_output.bpp], ignore_index=True)
+ bcm = pd.concat([self.output.bcm, self.singles_output.bcm], ignore_index=True)
+ initC = pd.concat([self.output.initC, self.singles_output.initC], ignore_index=True)
+ kick_info = pd.concat([self.output.kick_info, self.singles_output.kick_info], ignore_index=True)
+
+ return COSMICOutput(
+ bpp=bpp,
+ bcm=bcm,
+ initC=initC,
+ kick_info=kick_info,
+ label=self.label + ' [combined]' if self.label is not None else None
+ )
+
+
+def save_initC(filename, initC, key="initC", settings_key="initC_settings", force_save_all=False):
+ """Save an initC table to an HDF5 file.
+
+ Any column where every binary has the same value (setting) is saved separately with only a single copy
+ to save space.
+
+ This will take slightly longer (a few seconds instead of 1 second) to run but will save you around
+ a kilobyte per binary, which adds up!
+
+ Parameters
+ ----------
+ filename : `str`
+ Filename/path to the HDF5 file
+ initC : `pandas.DataFrame`
+ Initial conditions table
+ key : `str`, optional
+ Dataset key to use for main table, by default "initC"
+ settings_key : `str`, optional
+ Dataset key to use for settings table, by default "initC_settings"
+ force_save_all : `bool`, optional
+ If true, force all settings columns to be saved in the main table, by default False
+ """
+
+ # for each column, check if all values are the same
+ uniques = initC.nunique(axis=0)
+ compress_cols = [col for col in initC.columns if uniques[col] == 1]
+
+ if len(compress_cols) == 0 or force_save_all:
+ # nothing to compress, just save the whole table
+ initC.to_hdf(filename, key=key)
+ else:
+ # save the main table without the compressed columns
+ initC.drop(columns=compress_cols).to_hdf(filename, key=key)
+
+ # save the compressed columns separately
+ settings_df = pd.DataFrame([{col: initC[col].iloc[0] for col in compress_cols}])
+ settings_df.to_hdf(filename, key=settings_key)
+
+
+def load_initC(filename, key="initC", settings_key="initC_settings"):
+ """Load an initC table from an HDF5 file.
+
+ If settings were saved separately, they are merged back into the main table.
+
+ Parameters
+ ----------
+ filename : `str`
+ Filename/path to the HDF5 file
+ key : `str`, optional
+ Dataset key to use for main table, by default "initC"
+ settings_key : `str`, optional
+ Dataset key to use for settings table, by default "initC_settings"
+
+ Returns
+ -------
+ initC : `pandas.DataFrame`
+ Initial conditions table
+ """
+
+ with h5.File(filename, 'r') as f:
+ has_settings = settings_key in f.keys()
+
+ initC = pd.read_hdf(filename, key=key)
+
+ if has_settings:
+ settings_df = pd.read_hdf(filename, key=settings_key)
+ initC.loc[:, settings_df.columns] = settings_df.values[0]
+
+ return initC
+
+
diff --git a/src/cosmic/plotting.py b/src/cosmic/plotting.py
index 58ef4c71..4aca0c23 100644
--- a/src/cosmic/plotting.py
+++ b/src/cosmic/plotting.py
@@ -44,9 +44,6 @@
"evolve_and_plot",
]
-rsun_in_au = 215.0954
-day_in_year = 365.242
-
# Colors
primary_color = "C0"
secondary_color = "C1"
diff --git a/src/cosmic/utils.py b/src/cosmic/utils.py
index d0849bb6..aa4a72b7 100644
--- a/src/cosmic/utils.py
+++ b/src/cosmic/utils.py
@@ -61,8 +61,6 @@
"check_initial_conditions",
"convert_kstar_evol_type",
"parse_inifile",
- "save_initC",
- "load_initC",
"pop_write",
"a_from_p",
"p_from_a",
@@ -360,77 +358,6 @@ def conv_select(bcm_save, bpp_save, final_kstar_1, final_kstar_2, method, conv_l
return conv_save, conv_lims_bin_num
-def save_initC(filename, initC, key="initC", settings_key="initC_settings", force_save_all=False):
- """Save an initC table to an HDF5 file.
-
- Any column where every binary has the same value (setting) is saved separately with only a single copy
- to save space.
-
- This will take slightly longer (a few seconds instead of 1 second) to run but will save you around
- a kilobyte per binary, which adds up!
-
- Parameters
- ----------
- filename : `str`
- Filename/path to the HDF5 file
- initC : `pandas.DataFrame`
- Initial conditions table
- key : `str`, optional
- Dataset key to use for main table, by default "initC"
- settings_key : `str`, optional
- Dataset key to use for settings table, by default "initC_settings"
- force_save_all : `bool`, optional
- If true, force all settings columns to be saved in the main table, by default False
- """
-
- # for each column, check if all values are the same
- uniques = initC.nunique(axis=0)
- compress_cols = [col for col in initC.columns if uniques[col] == 1]
-
- if len(compress_cols) == 0 or force_save_all:
- # nothing to compress, just save the whole table
- initC.to_hdf(filename, key=key)
- else:
- # save the main table without the compressed columns
- initC.drop(columns=compress_cols).to_hdf(filename, key=key)
-
- # save the compressed columns separately
- settings_df = pd.DataFrame([{col: initC[col].iloc[0] for col in compress_cols}])
- settings_df.to_hdf(filename, key=settings_key)
-
-
-def load_initC(filename, key="initC", settings_key="initC_settings"):
- """Load an initC table from an HDF5 file.
-
- If settings were saved separately, they are merged back into the main table.
-
- Parameters
- ----------
- filename : `str`
- Filename/path to the HDF5 file
- key : `str`, optional
- Dataset key to use for main table, by default "initC"
- settings_key : `str`, optional
- Dataset key to use for settings table, by default "initC_settings"
-
- Returns
- -------
- initC : `pandas.DataFrame`
- Initial conditions table
- """
-
- with h5.File(filename, 'r') as f:
- has_settings = settings_key in f.keys()
-
- initC = pd.read_hdf(filename, key=key)
-
- if has_settings:
- settings_df = pd.read_hdf(filename, key=settings_key)
- initC.loc[:, settings_df.columns] = settings_df.values[0]
-
- return initC
-
-
def pop_write(
dat_store,
log_file,