From 46a18b61bc8c040556462cf7039f82f99b6c8ae3 Mon Sep 17 00:00:00 2001
From: Lorenzo Pompili <lorenzo.pompili@aei.mpg.de>
Date: Wed, 22 Nov 2023 14:58:56 +0100
Subject: [PATCH] ModeArray in settings

ModeArray should be used in settings instead of the previous mode_array.
We keep "mode_array" and disallow using both at the same time.

* added unit test for the settings
---
 pyseobnr/generate_waveform.py            | 18 ++++-
 pyseobnr/models/SEOBNRv5HM.py            |  4 +-
 test/regression_tests/test_SEOBNRv5HM.py |  3 +-
 test/regression_tests/test_settings.py   | 87 ++++++++++++++++++++++++
 4 files changed, 105 insertions(+), 7 deletions(-)
 create mode 100644 test/regression_tests/test_settings.py

diff --git a/pyseobnr/generate_waveform.py b/pyseobnr/generate_waveform.py
index 6a1cc2f..8affa92 100644
--- a/pyseobnr/generate_waveform.py
+++ b/pyseobnr/generate_waveform.py
@@ -171,7 +171,7 @@ class GenerateWaveform:
             deltaT: Time spacing, in seconds - Default: 1/2048 seconds
             f_max: Maximum frequency, in Hz - Default: 1024 Hz
             deltaF: Frequency spacing, in Hz - Default: 0.125
-            mode_array: Mode content (only positive must be specified, e.g [(2,2),(2,1)]) - Default: None (all modes)
+            ModeArray / mode_array: Mode content (only positive must be specified, e.g [(2,2),(2,1)]) - Default: None (all modes)
             approximant: 'SEOBNRv5HM' or 'SEOBNRv5PHM'. Default: 'SEOBNRv5HM'
 
         """
@@ -213,6 +213,7 @@ class GenerateWaveform:
             "f_ref": 20.0,
             "deltaT": 1.0 / 2048.0,
             "deltaF": 0.0,
+            "ModeArray": None,
             "mode_array": None,
             "approximant": "SEOBNRv5HM",
             "conditioning": 2,
@@ -330,6 +331,9 @@ class GenerateWaveform:
         if parameters["postadiabatic_type"] not in ["numeric", "analytic"]:
             raise ValueError("Unrecognised setting for dynamics postadiabatic type.")
 
+        if parameters["ModeArray"] is not None and parameters["mode_array"] is not None:
+            raise ValueError("Only one setting can be specified between ModeArray and mode_array.")
+
         self.parameters = parameters
 
     def generate_td_modes(self):
@@ -392,10 +396,14 @@ class GenerateWaveform:
                     ]
                 )
 
-        if self.parameters["mode_array"] != None:
+        if self.parameters["mode_array"] is not None:
             settings["return_modes"] = self.parameters[
                 "mode_array"
             ]  # Select mode array
+        if self.parameters["ModeArray"] is not None:
+            settings["return_modes"] = self.parameters[
+                "ModeArray"
+            ]  # Select mode array
 
         if "lmax_nyquist" in self.parameters:
             settings.update(lmax_nyquist=self.parameters["lmax_nyquist"])
@@ -504,10 +512,14 @@ class GenerateWaveform:
             if "r_size_input" in self.parameters:
                 settings.update(r_size_input=self.parameters["r_size_input"])
 
-            if self.parameters["mode_array"] != None:
+            if self.parameters["mode_array"] is not None:
                 settings["return_modes"] = self.parameters[
                     "mode_array"
                 ]  # Select mode array
+            if self.parameters["ModeArray"] is not None:
+                settings["return_modes"] = self.parameters[
+                    "ModeArray"
+                ]  # Select mode array
 
             if "lmax_nyquist" in self.parameters:
                 settings.update(lmax_nyquist=self.parameters["lmax_nyquist"])
diff --git a/pyseobnr/models/SEOBNRv5HM.py b/pyseobnr/models/SEOBNRv5HM.py
index ef63059..138c705 100644
--- a/pyseobnr/models/SEOBNRv5HM.py
+++ b/pyseobnr/models/SEOBNRv5HM.py
@@ -254,7 +254,7 @@ class SEOBNRv5HM_opt(Model):
         """
         Re-initialize all parameters to make sure everything is reset
         """
-        self.eob_pars = EOBParams(phys_pars, {}, mode_array=self.computed_modes)
+        self.eob_pars = EOBParams(phys_pars, {}, mode_array=list(self.computed_modes))
         self.eob_pars.flux_params.rho_initialized = False
         self.eob_pars.flux_params.prefixes = np.array(self.prefixes)
         self.eob_pars.flux_params.prefixes_abs = np.abs(
@@ -770,7 +770,7 @@ class SEOBNRv5PHM_opt(Model):
         """
         Re-initialize all parameters to make sure everything is reset
         """
-        self.eob_pars = EOBParams(phys_pars, {}, mode_array=self.computed_modes)
+        self.eob_pars = EOBParams(phys_pars, {}, mode_array=list(self.computed_modes))
         self.eob_pars.flux_params.rho_initialized = False
         self.eob_pars.flux_params.prefixes = np.array(self.prefixes)
         self.eob_pars.flux_params.prefixes_abs = np.abs(
diff --git a/test/regression_tests/test_SEOBNRv5HM.py b/test/regression_tests/test_SEOBNRv5HM.py
index 472fa38..af66ee6 100644
--- a/test/regression_tests/test_SEOBNRv5HM.py
+++ b/test/regression_tests/test_SEOBNRv5HM.py
@@ -8,9 +8,8 @@ It should be regenerated using regenerate_SEOBNRv5HM.py whenever
 understood changes to the model are made.
 """
 
-import pytest
 import numpy as np
-from pyseobnr.generate_waveform import generate_modes_opt, GenerateWaveform
+from pyseobnr.generate_waveform import GenerateWaveform
 import lal
 
 
diff --git a/test/regression_tests/test_settings.py b/test/regression_tests/test_settings.py
new file mode 100644
index 0000000..7044c43
--- /dev/null
+++ b/test/regression_tests/test_settings.py
@@ -0,0 +1,87 @@
+import lal
+import numpy as np
+import pytest
+
+from pyseobnr.generate_waveform import GenerateWaveform
+
+
+@pytest.fixture
+def basic_settings():
+    m1 = 50.0
+    m2 = 30.0
+    params_dict = {
+        "mass1": m1,
+        "mass2": m2,
+        "spin1x": 0,
+        "spin1y": 0,
+        "spin1z": 0.5,
+        "spin2x": 0,
+        "spin2y": 0,
+        "spin2z": 0.1,
+        "deltaT": 1 / 2048.0,
+        "deltaF": 0.125,
+        "f22_start": 0.0157 / ((m1 + m2) * np.pi * lal.MTSUN_SI),
+        "phi_ref": 0.0,
+        "distance": 1.0,
+        "inclination": np.pi / 3.0,
+        "f_max": 1024.0,
+        "approximant": "SEOBNRv5HM",
+        "postadiabatic": False,
+    }
+    return params_dict
+
+
+def test_mode_arrays_settings(basic_settings):
+    """Checks the behaviour wrt. mode_array setting"""
+
+    # cannot accept both
+    value_error = (
+        r"Only one setting can be specified between .?ModeArray.? and .?mode_array.?."
+    )
+
+    with pytest.raises(
+        ValueError,
+        match=value_error,
+    ):
+        _ = GenerateWaveform(basic_settings | {"mode_array": [], "ModeArray": []})
+
+    with pytest.raises(
+        ValueError,
+        match=value_error,
+    ):
+        _ = GenerateWaveform(
+            basic_settings
+            | {"approximant": "SEOBNRv5PHM", "mode_array": [], "ModeArray": []}
+        )
+
+    calls_to_check = (
+        "generate_td_modes",
+        "generate_td_polarizations",
+        "generate_fd_polarizations",
+    )
+
+    # works with list and tuples
+    for approximant in "SEOBNRv5HM", "SEOBNRv5PHM":
+        wfm_gen = GenerateWaveform(
+            basic_settings
+            | {"approximant": approximant, "mode_array": [(2, 2), (3, 3)]}
+        )
+
+        for func in calls_to_check:
+            _ = getattr(wfm_gen, func)()
+
+        wfm_gen = GenerateWaveform(
+            basic_settings
+            | {"approximant": approximant, "mode_array": ((2, 2), (3, 3))}
+        )
+
+        for func in calls_to_check:
+            _ = getattr(wfm_gen, func)()
+
+        # incorrect mode array yields an error
+        wfm_gen = GenerateWaveform(
+            basic_settings | {"approximant": approximant, "mode_array": [2, 2, 3, 3]}
+        )
+        for func in calls_to_check:
+            with pytest.raises(TypeError):
+                _ = getattr(wfm_gen, func)()
-- 
GitLab