PycWB 0.27.1__tar.gz → 0.28.0__tar.gz

{pycwb-0.27.1 → pycwb-0.28.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: PycWB
-Version: 0.27.1
+Version: 0.28.0
 Summary: A Python package for coherent gravitational wave burst analysis
 Home-page: https://git.ligo.org/yumeng.xu/pycwb
 Author-email: The PycWB team <yumeng.xu@ligo.org>

{pycwb-0.27.1 → pycwb-0.28.0}/PycWB.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: PycWB
-Version: 0.27.1
+Version: 0.28.0
 Summary: A Python package for coherent gravitational wave burst analysis
 Home-page: https://git.ligo.org/yumeng.xu/pycwb
 Author-email: The PycWB team <yumeng.xu@ligo.org>

{pycwb-0.27.1 → pycwb-0.28.0}/PycWB.egg-info/SOURCES.txt

@@ -15,6 +15,7 @@ PycWB.egg-info/top_level.txt
 benchmark/README.md
 benchmark/likelihood/README.md
 benchmark/likelihood/generate_data_for_likelihood.py
+benchmark/likelihood/perf_cuda.py
 benchmark/likelihood/performance_test_dpf.py
 benchmark/likelihood/performance_test_opt_sky.py
 benchmark/likelihood/user_parameters_injection.yaml

{pycwb-0.27.1 → pycwb-0.28.0}/benchmark/likelihood/generate_data_for_likelihood.py

@@ -3,9 +3,9 @@ import os
 from pycwb.config import Config
 from pycwb.modules.logger import logger_init
 
-if not os.environ.get('HOME_WAT_FILTERS'):
-    print('Please set HOME_WAT_FILTERS to the directory of WAT filters')
-    exit(1)
+# if not os.environ.get('HOME_WAT_FILTERS'):
+#     print('Please set HOME_WAT_FILTERS to the directory of WAT filters')
+#     exit(1)
 
 logger_init()
 
@@ -14,13 +14,21 @@ config = Config('./user_parameters_injection.yaml')
 #%% md
 ## generate injected data for each detector with given parameters in config
 #%%
-from pycwb.modules.read_data import generate_injection
-from pycwb.modules.job_segment import create_job_segment_from_injection
+from pycwb.modules.read_data import generate_injection, read_from_job_segment, generate_noise_for_job_seg
+from pycwb.modules.job_segment import create_job_segment_from_injection, create_job_segment_from_config
 
-job_segments = create_job_segment_from_injection(config.ifo, config.simulation, config.injection)
+job_segments = create_job_segment_from_config(config) #create_job_segment_from_injection(config.ifo, config.simulation, config.injection)
 
-data = generate_injection(config, job_segments[0])
+job_seg = job_segments[0]
 
+data = None
+# data = generate_injection(config, job_segments[0])
+if job_seg.frames:
+    data = read_from_job_segment(config, job_seg)
+if job_seg.noise:
+    data = generate_noise_for_job_seg(job_seg, config.inRate, data=data)
+if job_seg.injections:
+    data = generate_injection(config, job_seg, data)
 #%% md
 ## apply data conditioning to the data
 #%%

pycwb-0.28.0/benchmark/likelihood/perf_cuda.py

@@ -0,0 +1,168 @@
+import cupy as cp
+
+def dpf_cupy(FP, FX, rms):
+    """
+    Vectorized version of dpf_np using CuPy for GPU acceleration.
+    
+    :param FP: (n_sky, NIFO, 1) array
+    :param FX: (n_sky, NIFO, 1) array
+    :param rms: (NIFO, NPIX) array
+    :return: aa, fp, fx, si, co, ni (all as cupy arrays)
+    """
+    # Broadcast multiply
+    f = FP * rms[cp.newaxis, :, :]  # Shape: (n_sky, NIFO, NPIX)
+    F = FX * rms[cp.newaxis, :, :]
+
+    # Sum over NIFO (axis=1)
+    _ff = (f ** 2).sum(axis=1)
+    _FF = (F ** 2).sum(axis=1)
+    _fF = (f * F).sum(axis=1)
+
+    # Intermediate terms
+    _si = 2.0 * _fF
+    _co = _ff - _FF
+    _AP = _ff + _FF
+    _nn = cp.sqrt(_co**2 + _si**2)
+    _cc = _co / (_nn + 0.0001)
+    fp = (_AP + _nn) / 2.0
+
+    # Compute co and si with sign correction
+    si = cp.sqrt((1.0 - _cc) / 2.0)
+    co = cp.where(_si > 0.0, 
+                  cp.sqrt((1.0 + _cc) / 2.0), 
+                  -cp.sqrt((1.0 + _cc) / 2.0))
+
+    # Expand dimensions for broadcasting
+    co_exp = co[:, cp.newaxis, :]
+    si_exp = si[:, cp.newaxis, :]
+
+    # Update basis vectors
+    f_new = f * co_exp + F * si_exp
+    F_new = F * co_exp - f * si_exp
+
+    # Compute new correlation
+    fF_new = (f_new * F_new).sum(axis=1) / (fp + 0.0001)
+
+    # Adjust F
+    F_adj = F_new - f_new * fF_new[:, cp.newaxis, :]
+
+    # Final energy terms
+    fx = (F_adj ** 2).sum(axis=1)
+    ni = (f_new ** 4).sum(axis=1) / (fp**2 + 0.0001)
+
+    # Network statistics
+    NI = (fx / (ni + 0.0001)).sum(axis=1)
+    NN = (fp > 0.0).sum(axis=1)
+    aa = cp.sqrt(NI / (NN + 0.01))
+
+    return aa, fp, fx, si, co, ni
+
+
+def dpf_cupy(Fp0, Fx0, rms):
+    """
+    :param Fp0: Fp0 is a NIFO x 1 array (CuPy array)
+    :param Fx0: Fx0 is a NIFO x 1 array (CuPy array)
+    :param rms: rms is a NIFO x NPIX array (CuPy array)
+    :return: Tuple of results as CuPy arrays
+    """
+    f = rms * Fp0
+    F = rms * Fx0
+
+    _ff = cp.sum(f * f, axis=1)
+    _FF = cp.sum(F * F, axis=1)
+    _fF = cp.sum(f * F, axis=1)
+
+    _si = 2.0 * _fF
+    _co = _ff - _FF
+    _AP = _ff + _FF
+    _nn = cp.sqrt(_co**2 + _si**2)
+    _cc = _co / (_nn + 0.0001)
+    fp = (_AP + _nn) / 2.0
+
+    si = cp.sqrt((1.0 - _cc) / 2.0)
+    co = cp.where(_si > 0.0, 
+                 cp.sqrt((1.0 + _cc) / 2.0), 
+                 -cp.sqrt((1.0 + _cc) / 2.0))
+
+    # Update f and F with broadcasting
+    f = f * co[:, cp.newaxis] + F * si[:, cp.newaxis]
+    F = F * co[:, cp.newaxis] - f * si[:, cp.newaxis]
+
+    fF_new = cp.sum(f * F, axis=1) / (fp + 0.0001)
+    F -= f * fF_new[:, cp.newaxis]
+
+    fx = cp.sum(F * F, axis=1)
+    ni = cp.sum(f**4, axis=1) / (fp**2 + 0.0001)
+
+    NI = cp.sum(fx / (ni + 0.0001))
+    NN = cp.sum(fp > 0.0)
+
+    return cp.sqrt(NI / (NN + 0.01)), fp, fx, si, co, ni
+            
+
+def calculate_dpf_gpu(FP, FX, rms, n_sky: int, gamma_regulator: float, 
+                     network_energy_threshold: float):
+    """
+    GPU-accelerated version of calculate_dpf using CuPy.
+    
+    :param FP: (n_sky, NIFO, 1) array
+    :param FX: (n_sky, NIFO, 1) array
+    :param rms: (NIFO, NPIX) array
+    :return: network energy statistic
+    """
+    # Convert to float32 and ensure on GPU
+    FP = cp.asarray(FP, dtype=cp.float32)
+    FX = cp.asarray(FX, dtype=cp.float32)
+    rms = cp.asarray(rms, dtype=cp.float32)
+
+    # Compute all aa values in parallel
+    aa, _, _, _, _, _ = dpf_cupy(FP, FX, rms)
+
+    # Compute selection mask
+    mm = (aa > gamma_regulator).astype(cp.uint8)
+    ff = mm.sum()
+    FF = n_sky  # Original code sets all MM[i] = 1
+
+    # Final calculation
+    return (FF**2 / (ff**2 + 1e-9) - 1) * network_energy_threshold
+
+
+
+import pickle
+import numpy as np
+
+from pycwb.modules.likelihoodWP.likelihood import load_data_from_pixels
+
+# load FP, FX, rms, n_sky, gamma_regulator, network_energy_threshold from pickle
+with open('test_data.pkl', 'rb') as f:
+    test_data = pickle.load(f)
+
+FP = test_data['FP']
+FX = test_data['FX']
+pixels = test_data['pixels']
+n_ifo = test_data['n_ifo']
+n_sky = test_data['n_sky']
+gamma_regulator = test_data['gamma_regulator']
+network_energy_threshold = test_data['network_energy_threshold']
+rms, td00, td90, td_energy = load_data_from_pixels(pixels, n_ifo)
+
+
+calculate_dpf_gpu(FP.T, FX.T, rms.T,
+              int(n_sky),
+              gamma_regulator, network_energy_threshold)
+
+print(dpf_cupy(FP.T[1000].astype(np.float32), FX.T[1000].astype(np.float32), rms.T.astype(np.float32))[0])
+
+import time
+
+total_time = 0
+# # convert FP, FX, rms to float32
+for i in range(10):
+    start = time.time()
+    calculate_dpf_gpu(FP.T, FX.T, rms.T,
+                  int(n_sky),
+                  gamma_regulator, network_energy_threshold)
+    end = time.time()
+    total_time += end - start
+    print(end - start)
+print(f"Average time: {total_time / 10:.4f} s")

{pycwb-0.27.1 → pycwb-0.28.0}/pycwb/_version.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.27.1'
-__version_tuple__ = version_tuple = (0, 27, 1)
+__version__ = version = '0.28.0'
+__version_tuple__ = version_tuple = (0, 28, 0)

{pycwb-0.27.1 → pycwb-0.28.0}/pycwb/constants/user_parameters_schema.py

@@ -81,6 +81,12 @@ schema = {
             "default": False,
             "cwb": False
         },
+        "segment_processer": {
+            "type": "string",
+            "description": "main segment processer function for the analysis, the arguments for the function are: func(working_dir, config, job_seg, **kwargs)",
+            "default": "pycwb.workflow.subflow.process_job_segment.process_job_segment",
+            "cwb": False
+        },
         "use_root_supercluster": {
             "type": "boolean",
             "description": "use cwb supercluster, if false use python version",

{pycwb-0.27.1 → pycwb-0.28.0}/pycwb/modules/reconstruction/getMRAwaveform.py

@@ -8,7 +8,7 @@ from numba import njit
 logger = logging.getLogger(__name__)
 
 
-def get_MRA_wave(cluster, wdmList, rate, ifo, a_type, mode, nproc, whiten=False):
+def get_MRA_wave(cluster, wdmList, rate, ifo, a_type, mode, nproc, whiten=False) -> TimeSeries:
     """
     get MRA waveforms of type atype in time domain given lag nomber and cluster ID
 
@@ -122,9 +122,9 @@ def get_network_MRA_wave(config, cluster, rate, nIFO, rTDF, a_type, mode, tof, w
         # apply time delay
         if tof:
             R = rTDF  # effective time-delay rate
-            t_shift = v[i] / R
+            t_shift = -v[i] / R
             xf = x.to_frequencyseries()
-            xf = xf.cyclic_time_shift(t_shift)
+            xf.data *= np.exp(-2j * np.pi * xf.sample_frequencies * t_shift)
             x = xf.to_timeseries()
 
         waveforms.append(x)

{pycwb-0.27.1 → pycwb-0.28.0}/pycwb/utils/module.py

@@ -13,3 +13,13 @@ def import_helper(module_str, module_name):
         module = importlib.import_module(module_str)
 
     return module
+
+
+def import_function(func_str):
+    func_name = func_str.split('.')[-1]
+    module_name = '.'.join(func_str.split('.')[:-1])
+    if not module_name.startswith('pycwb'):
+        module_name = f"pycwb.modules.{module_name}"
+    module = import_helper(module_name, module_name)
+    func = getattr(module, func_name)
+    return func

{pycwb-0.27.1 → pycwb-0.28.0}/pycwb/workflow/batch.py

@@ -2,13 +2,13 @@ import multiprocessing
 import os
 import getpass
 import shutil
+import logging
 from pathlib import Path
-
-from htcondor import dags
-
 from pycwb.modules.logger import logger_init, log_prints
 from pycwb.workflow.subflow import prepare_job_runs, load_batch_run
-from pycwb.workflow.subflow.process_job_segment import process_job_segment
+from pycwb.utils.module import import_function
+
+logger = logging.getLogger(__name__)
 
 
 def batch_setup(file_name, working_dir='.',
@@ -17,6 +17,7 @@ def batch_setup(file_name, working_dir='.',
                 accounting_group=None, job_per_worker=10, n_proc=1, memory="6GB", disk="4GB",
                 dry_run=False, submit=False):
     import htcondor
+    from htcondor import dags
 
     logger_init(log_file, log_level)
     job_segments, config, working_dir = prepare_job_runs(working_dir, file_name, n_proc, dry_run, overwrite,
@@ -139,6 +140,10 @@ def batch_run(config_file, working_dir='.', log_file=None, log_level="INFO",
                                                        n_proc=n_proc, compress_json=compress_json)
     logger_init(log_file, log_level)
 
+    logger.info(f"Loading segment processer: {config.segment_processer}")
+    main_func = import_function(config.segment_processer)
+    logger.info(f"Segment processer loaded: {main_func}")
+
     exceptions = []
     for job_seg in job_segments:
         # check if the job is done
@@ -152,7 +157,7 @@ def batch_run(config_file, working_dir='.', log_file=None, log_level="INFO",
             #                                   args=(working_dir, config, job_seg, compress_json, catalog_file))
             # process.start()
             # process.join()
-            process_job_segment(working_dir, config, job_seg,
+            main_func(working_dir, config, job_seg,
                                 compress_json=compress_json, catalog_file=catalog_file)
 
             # create a flag file to indicate the job is done

{pycwb-0.27.1 → pycwb-0.28.0}/pycwb/workflow/run.py

@@ -1,6 +1,9 @@
 from pycwb.modules.logger import logger_init
 from pycwb.workflow.subflow import prepare_job_runs
-from pycwb.workflow.subflow.process_job_segment import process_job_segment
+from pycwb.utils.module import import_function
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 def search(file_name, working_dir='.', overwrite=False, log_file=None, log_level="INFO",
@@ -17,7 +20,11 @@ def search(file_name, working_dir='.', overwrite=False, log_file=None, log_level
     # cluster.scale(n_proc)
     # client = Client(cluster)
 
+    logger.info(f"Loading segment processer: {config.segment_processer}")
+    main_func = import_function(config.segment_processer)
+    logger.info(f"Segment processer loaded: {main_func}")
+
     for job_seg in job_segments:
-        process_job_segment(working_dir, config, job_seg, compress_json=compress_json)
+        main_func(working_dir, config, job_seg, compress_json=compress_json)
 
     # client.close()