From 0ccc5c9039d45a28cbf5bc8f8b4c8e1d9177af39 Mon Sep 17 00:00:00 2001
From: Daniel Williams <daniel.williams@ligo.org>
Date: Mon, 12 Aug 2024 12:26:04 +0000
Subject: [PATCH] Add psd file writing

---
 heron/injection.py            | 17 ++++++++++++-----
 heron/models/__init__.py      |  5 ++++-
 heron/models/lalnoise.py      | 10 ++++++++++
 tests/models/__init__.py      | 13 +++++++++++++
 tests/models/test_lalnoise.py | 10 ++++++++++
 tests/test_injection.py       |  1 +
 6 files changed, 50 insertions(+), 6 deletions(-)
 create mode 100644 tests/models/test_lalnoise.py

diff --git a/heron/injection.py b/heron/injection.py
index 234441e..27ebcb2 100644
--- a/heron/injection.py
+++ b/heron/injection.py
@@ -20,11 +20,12 @@
 
 
 def make_injection(
-    waveform=IMRPhenomPv2,
-    injection_parameters={},
-    times=None,
-    detectors=None,
-    framefile=None,
+        waveform=IMRPhenomPv2,
+        injection_parameters={},
+        times=None,
+        detectors=None,
+        framefile=None,
+        psdfile=None,
 ):
 
     parameters = {"ra": 0, "dec": 0, "psi": 0, "theta_jn": 0, "phase": 0}
@@ -60,6 +61,11 @@ def make_injection(
             logger.info(f"Saving framefile to {filename}")
             injection.write(filename, format="gwf")
 
+        if psdfile:
+            # Write the PSD file to an ascii file
+            filename = f"{detector.abbreviation}_{psdfile}.dat"
+            psd_model.to_file(filename)
+            
     return injections
 
 
@@ -147,5 +153,6 @@ def injection(settings):
         injection_parameters=parameters,
         detectors=detector_dict,
         framefile="injection",
+        psdfile="psd",
     )
     data = injections
diff --git a/heron/models/__init__.py b/heron/models/__init__.py
index e535f84..917fa10 100644
--- a/heron/models/__init__.py
+++ b/heron/models/__init__.py
@@ -1,4 +1,5 @@
 import torch
+import numpy as np
 from lal import antenna, MSUN_SI
 from astropy import units as u
 
@@ -74,8 +75,10 @@ class WaveformSurrogate(WaveformModel):
 
 
 class PSDModel:
-    pass
 
+    def to_file(self, filename, *args, **kwargs):
+        data = self.twocolumn(*args, **kwargs)
+        np.savetxt(filename, data)
 
 class PSDApproximant(PSDModel):
     pass
diff --git a/heron/models/lalnoise.py b/heron/models/lalnoise.py
index 528a569..0a5cc9e 100644
--- a/heron/models/lalnoise.py
+++ b/heron/models/lalnoise.py
@@ -42,6 +42,16 @@ def frequency_domain(
         psd = PSD(psd_data, frequencies=frequencies)
         return psd
 
+    def twocolumn(self, *args, **kwargs):
+        """
+        Produce the PSD in two-column format.
+        """
+        psd = self.frequency_domain(*args, **kwargs)
+        frequencies = psd.frequencies.value
+        data = np.array(psd.data)
+
+        return np.vstack([frequencies, data]).T
+    
     def covariance_matrix(self, times):
         """
         Return a time-domain representation of this power spectral density.
diff --git a/tests/models/__init__.py b/tests/models/__init__.py
index 453a678..e80b495 100644
--- a/tests/models/__init__.py
+++ b/tests/models/__init__.py
@@ -179,3 +179,16 @@ def test_waveform_mass_no_units(self):
                                                            "mass_2": mass2.to(u.kilogram).value})
 
         self.assertTrue(time_domain_1['plus'].data == time_domain_2['plus'].data)
+
+
+class _GenericPSD(unittest.TestCase):
+    
+    @classmethod
+    def setUpClass(self):
+        self.psd_model = AdvancedLIGO
+
+    def test_frequency_series_ascii_two_column(self):
+        """Test that a frequency series is correctly produced."""
+        frequency_series = self.psd_model().twocolumn()
+
+        self.assertTrue(frequency_series.ndim == 2)
diff --git a/tests/models/test_lalnoise.py b/tests/models/test_lalnoise.py
new file mode 100644
index 0000000..59e1399
--- /dev/null
+++ b/tests/models/test_lalnoise.py
@@ -0,0 +1,10 @@
+import unittest
+
+from heron.models.lalnoise import AdvancedLIGO
+from . import _GenericPSD
+
+class TestAdvancedLIGO(_GenericPSD):
+
+    @classmethod
+    def setUpClass(self):
+        self.psd_model = AdvancedLIGO
diff --git a/tests/test_injection.py b/tests/test_injection.py
index 91fe9d0..61ba28f 100644
--- a/tests/test_injection.py
+++ b/tests/test_injection.py
@@ -55,4 +55,5 @@ def test_multiple_inject(self):
                        detectors={"AdvancedLIGOHanford": "AdvancedLIGO",
                                   "AdvancedLIGOLivingston": "AdvancedLIGO"},
                        framefile="test",
+                       psdfile="psd"
                        )