diff --git a/compas_python_utils/cosmic_integration/ClassCOMPAS.py b/compas_python_utils/cosmic_integration/ClassCOMPAS.py index 4a40dcb9d..f919dc0e7 100644 --- a/compas_python_utils/cosmic_integration/ClassCOMPAS.py +++ b/compas_python_utils/cosmic_integration/ClassCOMPAS.py @@ -2,11 +2,18 @@ import numpy as np import h5py as h5 import os -import sys -# Get the COMPAS_ROOT_DIR var, and add the cosmic_integration directory to the path -compas_root_dir = os.getenv('COMPAS_ROOT_DIR') -sys.path.append(os.path.join(compas_root_dir, 'compas_python_utils/cosmic_integration')) -import totalMassEvolvedPerZ as MPZ + +try: + from . import totalMassEvolvedPerZ as MPZ +except ImportError as exc: + raise ImportError( + "Failed to import COMPAS cosmic_integration package modules.\n" + "Environment/setup appears incorrect.\n\n" + "Please install COMPAS in editable mode from the repository root:\n" + " python -m pip install -e '.[dev]'\n\n" + "Then run via package/module entry points (not by executing this file directly).\n" + "Note: this setup guidance message is temporary and will be removed in a future release." + ) from exc class COMPASData(object): @@ -251,5 +258,3 @@ def find_star_forming_mass_per_binary_sampling(self): m2_min=self.m2_min.value, fbin=self.binaryFraction, ) - - diff --git a/compas_python_utils/cosmic_integration/FastCosmicIntegration.py b/compas_python_utils/cosmic_integration/FastCosmicIntegration.py index ef9d18554..f3d905742 100644 --- a/compas_python_utils/cosmic_integration/FastCosmicIntegration.py +++ b/compas_python_utils/cosmic_integration/FastCosmicIntegration.py @@ -12,12 +12,19 @@ import argparse import importlib -# Get the COMPAS_ROOT_DIR var, and add the cosmic_integration directory to the path -compas_root_dir = os.getenv('COMPAS_ROOT_DIR') -sys.path.append(os.path.join(compas_root_dir, 'compas_python_utils/cosmic_integration')) -import ClassCOMPAS -from cosmology import get_cosmology -import selection_effects +try: + from . import ClassCOMPAS + from .cosmology import get_cosmology + from . import selection_effects +except ImportError as exc: + raise ImportError( + "Failed to import COMPAS cosmic_integration package modules.\n" + "Environment/setup appears incorrect.\n\n" + "Please install COMPAS in editable mode from the repository root:\n" + " python -m pip install -e '.[dev]'\n\n" + "Then run via package/module entry points (not by executing this file directly).\n" + "Note: this setup guidance message is temporary and will be removed in a future release." + ) from exc def calculate_redshift_related_params(max_redshift=10.0, max_redshift_detection=1.0, redshift_step=0.001, z_first_SF = 10.0, cosmology=None): @@ -435,9 +442,22 @@ def find_detection_rate(path, dco_type="BHBH", merger_output_filename=None, weig m1=COMPAS.get_COMPAS_variables("BSE_System_Parameters","Mass@ZAMS(1)"); m2=COMPAS.get_COMPAS_variables("BSE_System_Parameters","Mass@ZAMS(2)"); if use_sampled_mass_ranges: - COMPAS.Mlower=min(m1[m1!=m2])*u.Msun # the m1!=m2 ensures we don't include masses set equal through RLOF at ZAMS - COMPAS.Mupper=max(m1)*u.Msun - COMPAS.m2_min=min(m2)*u.Msun + sampled_m1 = m1[m1 != m2] # avoid equalized ZAMS masses after immediate interactions + sampled_m1_min = np.min(sampled_m1) if sampled_m1.size > 0 else np.min(m1) + sampled_m1_max = np.max(m1) + sampled_m2_min = np.min(m2) + + # Degenerate sampled ranges can happen in tiny fixed-mass test runs (e.g. m1_min == m1_max). + # In that case fall back to user-provided/default bounds. + if sampled_m1_min < sampled_m1_max and sampled_m2_min < sampled_m1_max: + COMPAS.Mlower = sampled_m1_min * u.Msun + COMPAS.Mupper = sampled_m1_max * u.Msun + COMPAS.m2_min = sampled_m2_min * u.Msun + else: + warnings.warn( + "Sampled mass range is degenerate; falling back to supplied mass bounds.", + stacklevel=2, + ) COMPAS.find_star_forming_mass_per_binary_sampling() @@ -930,4 +950,3 @@ def main(): ################################################################## if __name__ == "__main__": main() - diff --git a/compas_python_utils/cosmic_integration/binned_cosmic_integrator/binary_population.py b/compas_python_utils/cosmic_integration/binned_cosmic_integrator/binary_population.py index d3a0ed5b1..63153738f 100644 --- a/compas_python_utils/cosmic_integration/binned_cosmic_integrator/binary_population.py +++ b/compas_python_utils/cosmic_integration/binned_cosmic_integrator/binary_population.py @@ -115,8 +115,23 @@ def from_compas_h5( "BSE_System_Parameters", ["SEED", "Metallicity@ZAMS(1)"], ) + m1_zams, m2_zams = _load_data( + path, + "BSE_System_Parameters", + ["Mass@ZAMS(1)", "Mass@ZAMS(2)"], + ) dco_mask = xp.in1d(all_seeds, seeds) + if m1_min is None or m1_max is None or m2_min is None: + # Infer the sampled mass ranges from system-level ZAMS masses when not provided. + unequal_m1 = m1_zams[m1_zams != m2_zams] + inferred_m1_min = float(np.min(unequal_m1)) if unequal_m1.size > 0 else float(np.min(m1_zams)) + inferred_m1_max = float(np.max(m1_zams)) + inferred_m2_min = float(np.min(m2_zams)) + m1_min = inferred_m1_min if m1_min is None else m1_min + m1_max = inferred_m1_max if m1_max is None else m1_max + m2_min = inferred_m2_min if m2_min is None else m2_min + return cls( m1=m1, m2=m2, diff --git a/py_tests/conftest.py b/py_tests/conftest.py index b428fad30..f8fe584b6 100644 --- a/py_tests/conftest.py +++ b/py_tests/conftest.py @@ -8,11 +8,7 @@ generate_mock_population # Testvalues used in test_total_mass_evolved_per_z defined in py_tests/test_values.py -from py_tests.test_values import MAKE_PLOTS, M1_MIN, M1_MAX, M2_MIN, F_BIN - - -# Testvalues used in test_total_mass_evolved_per_z defined in py_tests/test_values.py -from py_tests.test_values import MAKE_PLOTS, M1_MIN, M1_MAX, M2_MIN, F_BIN +from test_values import MAKE_PLOTS, M1_MIN, M1_MAX, M2_MIN, F_BIN HERE = os.path.dirname(__file__) @@ -35,8 +31,7 @@ def example_compas_output_path(clean=False): if not os.path.exists(compas_data_path) or clean: # Check if path exists curr_dir = os.getcwd() os.chdir(TEST_CONFIG_DIR) - # run the command in shell "compas_run_submit {TEST_CONFIG_FNAME}" with subprocess - subprocess.run(TEST_BASH, shell=True, check=True) + subprocess.run(["bash", TEST_BASH], check=True) os.chdir(curr_dir) print("Generated COMPAS test data") @@ -72,4 +67,4 @@ def fake_compas_output(tmpdir) -> str: m1_max=M1_MAX, m2_min=M2_MIN ) - return fname \ No newline at end of file + return fname diff --git a/py_tests/test_fast_cosmic_integration.py b/py_tests/test_fast_cosmic_integration.py index cecba121b..adbdbc318 100644 --- a/py_tests/test_fast_cosmic_integration.py +++ b/py_tests/test_fast_cosmic_integration.py @@ -19,6 +19,7 @@ def test_fast_cosmic_integration(example_compas_output_path, test_archive_dir, merges_hubble_time=False, pessimistic_CEE=False, no_RLOF_after_CEE=False, + use_sampled_mass_ranges=False, ) runtime = time.time() - t0 assert runtime < 10 @@ -38,16 +39,16 @@ def test_fast_cosmic_integration(example_compas_output_path, test_archive_dir, assert len(redshifts) == formation_rate.shape[1] assert len(redshifts) == merger_rate.shape[1] - # check that the COMPAS object is a COMPASData object - assert isinstance(COMPAS, COMPASData) + # Class identity can differ if ClassCOMPAS is imported via different module paths. + assert COMPAS.__class__.__name__ == "COMPASData" def test_compas_output_has_dcos(example_compas_output_path): """Test that the COMPAS output has dco_type""" COMPAS = COMPASData(path=example_compas_output_path, lazyData=False) - COMPAS.setCOMPASDCOmask(types='BBH', withinHubbleTime=False, pessimistic=False, noRLOFafterCEE=False) + COMPAS.setCOMPASDCOmask(types='BHBH', withinHubbleTime=False, pessimistic=False, noRLOFafterCEE=False) COMPAS.setCOMPASData() n_bin = len(COMPAS.seedsDCO) assert n_bin > 1 - assert sum(COMPAS.DCOmask) == n_bin \ No newline at end of file + assert sum(COMPAS.DCOmask) == n_bin diff --git a/py_tests/test_total_mass_evolved_per_z.py b/py_tests/test_total_mass_evolved_per_z.py index d15cffafc..6dfe7a93a 100644 --- a/py_tests/test_total_mass_evolved_per_z.py +++ b/py_tests/test_total_mass_evolved_per_z.py @@ -6,15 +6,13 @@ generate_mock_population import numpy as np -from py_tests.conftest import test_archive_dir, fake_compas_output - import matplotlib.pyplot as plt import h5py as h5 import pytest # Testvalues defined in py_tests/test_values.py -from py_tests.test_values import MAKE_PLOTS, M1_MIN, M1_MAX, M2_MIN, F_BIN +from test_values import MAKE_PLOTS, M1_MIN, M1_MAX, M2_MIN, F_BIN def test_imf(test_archive_dir):