<jupyter_start><jupyter_text>This is a Moment Network example with gravitational wave signals Using simulated gravitational wave data to estimate the moments of the marginal posterior for the strain $h_+$ at each time step (these would form a $>100$-dimensional posterior distribution)  Here we use a dense neural network as a demo, though a CNN would work well for this data and can be easily used in the Moment Network hierarchy.<jupyter_code>import sys
!{sys.executable} -m pip install pycbc lalsuite ligo-common --no-cache-dir
!pip install 'git+https://github.com/NiallJeffrey/MomentNetworks.git'
import matplotlib.pyplot as plt
import numpy as np
import time
from scipy import signal
from google.colab import files
from pycbc.waveform import get_td_waveform
import pycbc.noise
import pycbc.psd
import random
from momentnetworks import demo
import gc<jupyter_output><empty_output><jupyter_text>Generate simulated waveforms and noisy data<jupyter_code>%matplotlib inline
n_training = 5000
n_training_samples = n_training*16
print('Number of training samples = '+str(n_training_samples))
mass_1_samples = np.random.uniform(10,30,n_training)
mass_2_samples = np.random.uniform(10,30,n_training)
distance_samples = np.random.uniform(500,1500,n_training)
sim_data, sim_data_noisy = demo.grav_wave_series(n_training,
                                                 mass_1_samples=mass_1_samples,
                                                 mass_2_samples=mass_2_samples,
                                                 distance_samples=distance_samples)
noise_alone = sim_data_noisy - sim_data
_ =plt.figure(figsize=(14,4))
_ = plt.subplot(1,2,1)
_ = plt.ylabel(r'$h_+$ strain',fontsize=14)
_ =plt.plot(np.linspace(0,1/16.,128)-1/8.,sim_data_noisy[63,:], alpha=0.5), plt.ylim(-.8,0.8)
_ =plt.plot(np.linspace(0,1/16.,128)-1/8.,sim_data[63,:], 'C1--'), plt.xlabel(r'$t\ /$ seconds',fontsize=14)
_ = plt.subplot(1,2,2)
_ = plt.ylabel(r'$h_+$ strain',fontsize=14)
_ =plt.plot(np.linspace(0,1/16.,128)-1/16.,sim_data_noisy[63+int(n_training*8),:], alpha=0.5), plt.ylim(-0.8,0.8)
_ =plt.plot(np.linspace(0,1/16.,128)-1/16.,sim_data[63+int(n_training*8),:], 'C1--'), plt.xlabel(r'$t\ /$ seconds',fontsize=14)

# shuffling is important, otherwise the mergers are all in the second half of the array
shuffled_indices = np.arange(sim_data_noisy.shape[0])
np.random.shuffle(shuffled_indices)
sim_data_noisy= sim_data_noisy[shuffled_indices]
sim_data= sim_data[shuffled_indices]
_ =plt.figure(figsize=(14,4))
_ = plt.subplot(1,2,1)
_ = plt.ylabel(r'$h_+$ strain',fontsize=14)
_ =plt.plot(np.linspace(0,1/16.,128)-1/8.,sim_data_noisy[63,:], alpha=0.5), plt.ylim(-.8,0.8)
_ =plt.plot(np.linspace(0,1/16.,128)-1/8.,sim_data[63,:], 'C1--'), plt.xlabel(r'$t\ /$ seconds',fontsize=14)
_ = plt.subplot(1,2,2)
_ = plt.ylabel(r'$h_+$ strain',fontsize=14)
_ =plt.plot(np.linspace(0,1/16.,128)-1/16.,sim_data_noisy[63+int(n_training*8),:], alpha=0.5), plt.ylim(-0.8,0.8)
_ =plt.plot(np.linspace(0,1/16.,128)-1/16.,sim_data[63+int(n_training*8),:], 'C1--'), plt.xlabel(r'$t\ /$ seconds',fontsize=14)<jupyter_output><empty_output><jupyter_text>Train moment network As long as the validation loss is decreasing, the estimate of the moment of the marginal posterior is improving<jupyter_code>gc.collect()
validation_index = np.int(n_training * 7)
model_instance = demo.simple_leaky(128,128, learning_rate=1e-4)
regression = model_instance.model()

history = regression.fit(sim_data_noisy[:validation_index]+0.5,
                               sim_data[:validation_index]+0.5,
                               epochs =500,batch_size =512, shuffle=True,
                               validation_data = (sim_data_noisy[validation_index:]+0.5,
                               sim_data[validation_index:]+0.5))
predicted_mean = regression.predict(np.atleast_2d(sim_data_noisy[:]+0.5))
training_var_unknown_mean = (sim_data[:]+0.5-predicted_mean)**2.
training_var_unknown_mean.shape
model_instance = demo.simpler_leaky(128,128, learning_rate=1e-3)
regression_var_unknown_mean = model_instance.model()
history = regression_var_unknown_mean.fit(sim_data_noisy[:validation_index]+0.5,
                                          training_var_unknown_mean[:validation_index]*1e3,
                                          epochs =100,batch_size =512, shuffle=True,
                                          validation_data = (sim_data_noisy[validation_index:]+0.5,
                                                             training_var_unknown_mean[validation_index:]*1e3))
for i_index in np.arange(57+int(n_training*8),63+int(n_training*8)):
  print(i_index)
  predicted_var_obs = (regression_var_unknown_mean.predict(np.atleast_2d(sim_data_noisy[i_index]+0.5))[0])*1e-3
  predicted_mean_obs = (regression.predict(np.atleast_2d(sim_data_noisy[i_index,]+0.5))[0])
  plt.figure(figsize=(12,4))
  _ = plt.subplot(1,2,1)
  _ = plt.ylabel(r'$h_+$ strain',fontsize=14)
  plt.fill_between(np.linspace(0,1./16,128),
                  predicted_mean_obs-0.5- np.sqrt(predicted_var_obs),
                  predicted_mean_obs + np.sqrt(predicted_var_obs)-0.5, alpha=0.5)
  plt.scatter(np.linspace(0,1./16,128),sim_data_noisy[i_index,], alpha=0.8)
  plt.plot(np.linspace(0,1./16,128),sim_data[i_index,], c='C2')
  plt.xlabel(r'$t\ /$ seconds',fontsize=14), plt.show()<jupyter_output>40057
{'fraction_non_alphanumeric': 0.14349859681945742, 'fraction_numerical': 0.055940130963517305, 'mean_word_length': 4.010309278350515, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 13, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '420767', 'license': 'MIT', 'n_tokens_mistral': 2063, 'n_tokens_neox': 1964, 'n_words': 292, 'path': 'MomentNetwork_demo/MomentNetwork_grav_wave_demo.ipynb', 'repo_name': 'Pablo-Lemos/MomentNetworks', 'size': '5345'}