Update Run Manager for New Infrastructure

example/Single_event_runManager.py

@@ -1,4 +1,3 @@
-
 import jax
 import jax.numpy as jnp
 
@@ -12,58 +11,50 @@
 mass_matrix = mass_matrix.at[5, 5].set(1e-3)
 mass_matrix = mass_matrix * 3e-3
 local_sampler_arg = {"step_size": mass_matrix}
-bounds = jnp.array(
-    [
-        [10.0, 40.0],
-        [0.125, 1.0],
-        [-1.0, 1.0],
-        [-1.0, 1.0],
-        [0.0, 2000.0],
-        [-0.05, 0.05],
-        [0.0, 2 * jnp.pi],
-        [-1.0, 1.0],
-        [0.0, jnp.pi],
-        [0.0, 2 * jnp.pi],
-        [-1.0, 1.0],
-    ]
-)
-
 
 run = SingleEventRun(
     seed=0,
     detectors=["H1", "L1"],
+    data_parameters={
+        "trigger_time": 1126259462.4,
+        "duration": 4,
+        "post_trigger_duration": 2,
+        "f_min": 20.0,
+        "f_max": 1024.0,
+        "tukey_alpha": 0.2,
+        "f_sampling": 4096.0,
+    },
     priors={
-        "M_c": {"name": "Uniform", "xmin": 10.0, "xmax": 80.0},
-        "q": {"name": "MassRatio"},
-        "s1_z": {"name": "Uniform", "xmin": -1.0, "xmax": 1.0},
-        "s2_z": {"name": "Uniform", "xmin": -1.0, "xmax": 1.0},
-        "d_L": {"name": "Uniform", "xmin": 0.0, "xmax": 2000.0},
-        "t_c": {"name": "Uniform", "xmin": -0.05, "xmax": 0.05},
-        "phase_c": {"name": "Uniform", "xmin": 0.0, "xmax": 2 * jnp.pi},
-        "cos_iota": {"name": "CosIota"},
-        "psi": {"name": "Uniform", "xmin": 0.0, "xmax": jnp.pi},
-        "ra": {"name": "Uniform", "xmin": 0.0, "xmax": 2 * jnp.pi},
-        "sin_dec": {"name": "SinDec"},
+        "M_c": {"name": "UniformPrior", "xmin": 10.0, "xmax": 80.0},
+        "q": {"name": "UniformPrior", "xmin": 0.0, "xmax": 1.0},
+        "s1_z": {"name": "UniformPrior", "xmin": -1.0, "xmax": 1.0},
+        "s2_z": {"name": "UniformPrior", "xmin": -1.0, "xmax": 1.0},
+        "d_L": {"name": "UniformPrior", "xmin": 1.0, "xmax": 2000.0},
+        "t_c": {"name": "UniformPrior", "xmin": -0.05, "xmax": 0.05},
+        "phase_c": {"name": "UniformPrior", "xmin": 0.0, "xmax": 2 * jnp.pi},
+        "iota": {"name": "SinePrior"},
+        "psi": {"name": "UniformPrior", "xmin": 0.0, "xmax": jnp.pi},
+        "ra": {"name": "UniformPrior", "xmin": 0.0, "xmax": 2 * jnp.pi},
+        "dec": {"name": "CosinePrior"},
     },
     waveform_parameters={"name": "RippleIMRPhenomD", "f_ref": 20.0},
-    jim_parameters={
-        "n_loop_training": 10,
-        "n_loop_production": 10,
-        "n_local_steps": 150,
-        "n_global_steps": 150,
-        "n_chains": 500,
-        "n_epochs": 50,
-        "learning_rate": 0.001,
-        "n_max_examples": 45000,
-        "momentum": 0.9,
-        "batch_size": 50000,
-        "use_global": True,
-        "keep_quantile": 0.0,
-        "train_thinning": 1,
-        "output_thinning": 10,
-        "local_sampler_arg": local_sampler_arg,
-    },
-    likelihood_parameters={"name": "HeterodynedTransientLikelihoodFD", "bounds": bounds},
+    likelihood_parameters={"name": "TransientLikelihoodFD"},
+    sample_transforms=[
+        {"name": "BoundToUnbound", "name_mapping": [["M_c"], ["M_c_unbounded"]], "original_lower_bound": 10.0, "original_upper_bound": 80.0,},
+        {"name": "BoundToUnbound", "name_mapping": [["q"], ["q_unbounded"]], "original_lower_bound": 0.0, "original_upper_bound": 1.0,},
+        {"name": "BoundToUnbound", "name_mapping": [["s1_z"], ["s1_z_unbounded"]], "original_lower_bound": -1.0, "original_upper_bound": 1.0,},
+        {"name": "BoundToUnbound", "name_mapping": [["s2_z"], ["s2_z_unbounded"]], "original_lower_bound": -1.0, "original_upper_bound": 1.0,},
+        {"name": "BoundToUnbound", "name_mapping": [["d_L"], ["d_L_unbounded"]], "original_lower_bound": 1.0, "original_upper_bound": 2000.0,},
+        {"name": "BoundToUnbound", "name_mapping": [["t_c"], ["t_c_unbounded"]], "original_lower_bound": -0.05, "original_upper_bound": 0.05,},
+        {"name": "BoundToUnbound", "name_mapping": [["phase_c"], ["phase_c_unbounded"]], "original_lower_bound": 0.0, "original_upper_bound": 2 * jnp.pi,},
+        {"name": "BoundToUnbound", "name_mapping": [["iota"], ["iota_unbounded"]], "original_lower_bound": 0.0, "original_upper_bound": jnp.pi,},
+        {"name": "BoundToUnbound", "name_mapping": [["psi"], ["psi_unbounded"]], "original_lower_bound": 0.0, "original_upper_bound": jnp.pi,},
+        {"name": "BoundToUnbound", "name_mapping": [["ra"], ["ra_unbounded"]], "original_lower_bound": 0.0, "original_upper_bound": 2 * jnp.pi,},
+        {"name": "BoundToUnbound", "name_mapping": [["dec"], ["dec_unbounded"]], "original_lower_bound": 0.0, "original_upper_bound": jnp.pi,},
+    ],
+    likelihood_transforms=[
+        {"name": "MassRatioToSymmetricMassRatioTransform", "name_mapping": [["q"], ["eta"]]},
+    ],
     injection=True,
     injection_parameters={
         "M_c": 28.6,
@@ -78,15 +69,28 @@
         "ra": 1.2,
         "dec": 0.3,
     },
-    data_parameters={
-        "trigger_time": 1126259462.4,
-        "duration": 4,
-        "post_trigger_duration": 2,
-        "f_min": 20.0,
-        "f_max": 1024.0,
-        "tukey_alpha": 0.2,
-        "f_sampling": 4096.0,
+    jim_parameters={
+        "n_loop_training": 100,
+        "n_loop_production": 20,
+        "n_local_steps": 10,
+        "n_global_steps": 1000,
+        "n_chains": 500,
+        "n_epochs": 30,
+        "learning_rate": 1e-4,
+        "n_max_examples": 30000,
+        "momentum": 0.9,
+        "batch_size": 30000,
+        "use_global": True,
+        "train_thinning": 1,
+        "output_thinning": 10,
+        "local_sampler_arg": local_sampler_arg,
     },
 )
 
 run_manager = SingleEventPERunManager(run=run)
+run_manager.sample()
+
+# plot the corner plot and diagnostic plot
+run_manager.plot_corner()
+run_manager.plot_diagnostic()
+run_manager.save_summary()

src/jimgw/single_event/detector.py

@@ -373,7 +373,7 @@ def inject_signal(
         h_sky: dict[str, Float[Array, " n_sample"]],
         params: dict[str, Float],
         psd_file: str = "",
-    ) -> None:
+    ) -> tuple[Float, Float]:
         """
         Inject a signal into the detector data.
 
@@ -392,7 +392,7 @@ def inject_signal(
 
         Returns
         -------
-        None
+        SNR
         """
         self.frequencies = freqs
         self.psd = self.load_psd(freqs, psd_file)
@@ -415,6 +415,8 @@ def inject_signal(
         print(f"The injected optimal SNR is {optimal_SNR}")
         print(f"The injected match filter SNR is {match_filter_SNR}")
 
+        return optimal_SNR, match_filter_SNR
+
     @jaxtyped(typechecker=typechecker)
     def load_psd(
         self, freqs: Float[Array, " n_sample"], psd_file: str = ""