From 7ba1ab3d9f983f25f2f654e32b9655f29d109bb1 Mon Sep 17 00:00:00 2001 From: "guillaume.grolleron" Date: Fri, 15 Sep 2023 16:31:36 +0200 Subject: [PATCH] docstrings container + unit test --- src/nectarchain/data/container/__init__.py | 3 +- .../data/container/chargesContainer.py | 233 +++++++++------- src/nectarchain/data/container/core.py | 1 + .../data/container/tests/test_charge.py | 252 ++++-------------- .../data/container/tests/test_waveforms.py | 76 ++++++ .../data/container/waveformsContainer.py | 68 ++++- 6 files changed, 330 insertions(+), 303 deletions(-) diff --git a/src/nectarchain/data/container/__init__.py b/src/nectarchain/data/container/__init__.py index 905aff5a..9df719d0 100644 --- a/src/nectarchain/data/container/__init__.py +++ b/src/nectarchain/data/container/__init__.py @@ -1,2 +1,3 @@ from .waveformsContainer import * -from .chargesContainer import * \ No newline at end of file +from .chargesContainer import * +from .core import * \ No newline at end of file diff --git a/src/nectarchain/data/container/chargesContainer.py b/src/nectarchain/data/container/chargesContainer.py index df7fb1af..d23634f4 100644 --- a/src/nectarchain/data/container/chargesContainer.py +++ b/src/nectarchain/data/container/chargesContainer.py @@ -4,7 +4,7 @@ log.handlers = logging.getLogger('__main__').handlers import numpy as np -from ctapipe.containers import Field +from ctapipe.container import Field from abc import ABC import os from pathlib import Path @@ -13,129 +13,172 @@ from .core import ArrayDataContainer class ChargesContainer(ArrayDataContainer): - - charges_hg = Field( - type = np.ndarray, - description = 'The high gain charges') - charges_lg = Field( - type = np.ndarray, - description = 'The low gain charges') - peak_hg = Field( - type = np.ndarray, - description = 'The high gain peak time') - peak_lg = Field( - type = np.ndarray, - description = 'The low gain peak time') - method = Field( - type = str, - description = 'The charge extraction method used') - + """ + A container that holds information about charges from a specific run. + + Fields: + charges_hg (np.ndarray): An array of high gain charges. + charges_lg (np.ndarray): An array of low gain charges. + peak_hg (np.ndarray): An array of high gain peak time. + peak_lg (np.ndarray): An array of low gain peak time. + method (str): The charge extraction method used. + """ + + charges_hg = Field( + type=np.ndarray(), + description='The high gain charges' + ) + charges_lg = Field( + type=np.ndarray, + description='The low gain charges' + ) + peak_hg = Field( + type=np.ndarray, + description='The high gain peak time' + ) + peak_lg = Field( + type=np.ndarray, + description='The low gain peak time' + ) + method = Field( + type=str, + description='The charge extraction method used' + ) class ChargesContainerIO(ABC) : - def write(path : Path, containers : ChargesContainer,**kwargs) : - """method to write in an output FITS file the ChargeContainer. - + """ + The `ChargesContainerIO` class provides methods for writing and loading `ChargesContainer` instances to/from FITS files. + + Example Usage: + # Writing a ChargesContainer instance to a FITS file + chargesContainer = ChargesContainer() + ChargesContainerIO.write(path, chargesContainer) + + # Loading a ChargesContainer instance from a FITS file + chargesContainer = ChargesContainerIO.load(path, run_number) + + Main functionalities: + - Writing a `ChargesContainer` instance to a FITS file. + - Loading a `ChargesContainer` instance from a FITS file. + + Methods: + - `write(path: Path, container: ChargesContainer, **kwargs) -> None`: Writes a `ChargesContainer` instance to a FITS file. The method takes a file path and the `ChargesContainer` instance as input. Additional keyword arguments can be provided to customize the file name and overwrite behavior. + - `load(path: Path, run_number: int, **kwargs) -> ChargesContainer`: Loads a `ChargesContainer` instance from a FITS file. The method takes a file path and the run number as input. Additional keyword arguments can be provided to specify an explicit file name. + + Fields: + The `ChargesContainerIO` class does not have any fields. + """ + @staticmethod + def write(path : Path, container : ChargesContainer,**kwargs) -> None: + """Write a ChargesContainer instance to a FITS file. Args: - path (str): the directory where you want to save data + path (str): The directory where the FITS file will be saved. + container (ChargesContainer): The ChargesContainer instance to be written to the FITS file. + **kwargs: Additional keyword arguments for customization. + Keyword Args: + suffix (str): A suffix to be added to the file name (default: ""). + overwrite (bool): Whether to overwrite the file if it already exists (default: False). + Returns: + None: This method does not return any value. + Raises: + OSError: If there is an error while writing the FITS file. + Exception: If there is any other exception during the writing process. + Example: + chargesContainer = ChargesContainer() + ChargesContainerIO.write(path, chargesContainer, suffix="v1", overwrite=True) """ suffix = kwargs.get("suffix","") if suffix != "" : suffix = f"_{suffix}" - log.info(f"saving in {path}") os.makedirs(path,exist_ok = True) - hdr = fits.Header() - hdr['RUN'] = containers.run_number - hdr['NEVENTS'] = containers.nevents - hdr['NPIXELS'] = containers.npixels - hdr['METHOD'] = containers.method - hdr['CAMERA'] = containers.camera + hdr['RUN'] = container.run_number + hdr['NEVENTS'] = container.nevents + hdr['NPIXELS'] = container.npixels + hdr['METHOD'] = container.method + hdr['CAMERA'] = container.camera + hdr['COMMENT'] = f"The charge containeur for run {container.run_number} with {container.method} method : primary is the pixels id, then you can find HG charge, LG charge, HG peak and LG peak, 2 last HDU are composed of event properties and trigger patern" - - hdr['COMMENT'] = f"The charge containeur for run {containers.run_number} with {containers.method} method : primary is the pixels id, then you can find HG charge, LG charge, HG peak and LG peak, 2 last HDU are composed of event properties and trigger patern" - - primary_hdu = fits.PrimaryHDU(containers.pixels_id,header=hdr) - charge_hg_hdu = fits.ImageHDU(containers.charges_hg,name = "HG charge") - charge_lg_hdu = fits.ImageHDU(containers.charges_lg,name = "LG charge") - peak_hg_hdu = fits.ImageHDU(containers.peak_hg, name = 'HG peak time') - peak_lg_hdu = fits.ImageHDU(containers.peak_lg, name = 'LG peak time') - - col1 = fits.Column(array = containers.broken_pixels_hg, name = "HG broken pixels", format = f'{containers.broken_pixels_hg.shape[1]}L') - col2 = fits.Column(array = containers.broken_pixels_lg, name = "LG broken pixels", format = f'{containers.broken_pixels_lg.shape[1]}L') + primary_hdu = fits.PrimaryHDU(container.pixels_id,header=hdr) + charge_hg_hdu = fits.ImageHDU(container.charges_hg,name = "HG charge") + charge_lg_hdu = fits.ImageHDU(container.charges_lg,name = "LG charge") + peak_hg_hdu = fits.ImageHDU(container.peak_hg, name = 'HG peak time') + peak_lg_hdu = fits.ImageHDU(container.peak_lg, name = 'LG peak time') + + col1 = fits.Column(array = container.broken_pixels_hg, name = "HG broken pixels", format = f'{container.broken_pixels_hg.shape[1]}L') + col2 = fits.Column(array = container.broken_pixels_lg, name = "LG broken pixels", format = f'{container.broken_pixels_lg.shape[1]}L') coldefs = fits.ColDefs([col1, col2]) broken_pixels = fits.BinTableHDU.from_columns(coldefs,name = 'trigger patern') - col1 = fits.Column(array = containers.event_id, name = "event_id", format = '1I') - col2 = fits.Column(array = containers.event_type, name = "event_type", format = '1I') - col3 = fits.Column(array = containers.ucts_timestamp, name = "ucts_timestamp", format = '1K') - col4 = fits.Column(array = containers.ucts_busy_counter, name = "ucts_busy_counter", format = '1I') - col5 = fits.Column(array = containers.ucts_event_counter, name = "ucts_event_counter", format = '1I') - col6 = fits.Column(array = containers.multiplicity, name = "multiplicity", format = '1I') - + col1 = fits.Column(array = container.event_id, name = "event_id", format = '1I') + col2 = fits.Column(array = container.event_type, name = "event_type", format = '1I') + col3 = fits.Column(array = container.ucts_timestamp, name = "ucts_timestamp", format = '1K') + col4 = fits.Column(array = container.ucts_busy_counter, name = "ucts_busy_counter", format = '1I') + col5 = fits.Column(array = container.ucts_event_counter, name = "ucts_event_counter", format = '1I') + col6 = fits.Column(array = container.multiplicity, name = "multiplicity", format = '1I') coldefs = fits.ColDefs([col1, col2, col3, col4, col5, col6]) event_properties = fits.BinTableHDU.from_columns(coldefs, name = 'event properties') - - col1 = fits.Column(array = containers.trig_pattern_all, name = "trig_pattern_all", format = f'{4 * containers.trig_pattern_all.shape[1]}L',dim = f'({ containers.trig_pattern_all.shape[1]},4)') - col2 = fits.Column(array = containers.trig_pattern, name = "trig_pattern", format = f'{containers.trig_pattern_all.shape[1]}L') + + col1 = fits.Column(array = container.trig_pattern_all, name = "trig_pattern_all", format = f'{4 * container.trig_pattern_all.shape[1]}L',dim = f'({ container.trig_pattern_all.shape[1]},4)') + col2 = fits.Column(array = container.trig_pattern, name = "trig_pattern", format = f'{container.trig_pattern_all.shape[1]}L') coldefs = fits.ColDefs([col1, col2]) trigger_patern = fits.BinTableHDU.from_columns(coldefs, name = 'trigger patern') - + hdul = fits.HDUList([primary_hdu, charge_hg_hdu, charge_lg_hdu,peak_hg_hdu,peak_lg_hdu,broken_pixels,event_properties,trigger_patern]) try : - hdul.writeto(Path(path)/f"charge_run{containers.run_number}{suffix}.fits",overwrite=kwargs.get('overwrite',False)) - log.info(f"charge saved in {Path(path)}/charge_run{containers.run_number}{suffix}.fits") + hdul.writeto(Path(path)/f"charge_run{container.run_number}{suffix}.fits",overwrite=kwargs.get('overwrite',False)) + log.info(f"charge saved in {Path(path)}/charge_run{container.run_number}{suffix}.fits") except OSError as e : log.warning(e) except Exception as e : log.error(e,exc_info = True) raise e - def load(path : Path,run_number : int,**kwargs) : - """load ChargeContainer from FITS file previously written with ChargeContainer.write() method - + @staticmethod + def load(path: str, run_number: int, **kwargs) -> ChargesContainer: + """Load a ChargesContainer instance from a FITS file. + This method opens a FITS file and retrieves the necessary data to create a ChargesContainer instance. + The FITS file should have been previously written using the write method of the ChargesContainerIO class. Args: - path (str): path of the FITS file - run_number (int) : the run number - + path (str): The path of the FITS file. + run_number (int): The run number. + **kwargs: Additional keyword arguments. + explicit_filename (str): If provided, the explicit filename to load. Returns: - ChargeContainer: ChargeContainer instance + ChargesContainer: The loaded ChargesContainer instance. + Example: + chargesContainer = ChargesContainerIO.load(path, run_number) """ - if kwargs.get("explicit_filename",False) : + if kwargs.get("explicit_filename", False): filename = kwargs.get("explicit_filename") - log.info(f"loading {filename}") - else : - log.info(f"loading in {path} run number {run_number}") - filename = Path(path)/f"charge_run{run_number}.fits" - - with fits.open(filename) as hdul : - containers = ChargesContainer() - containers.run_number = hdul[0].header['RUN'] - containers.nevents = hdul[0].header['NEVENTS'] - containers.npixels = hdul[0].header['NPIXELS'] - containers.method = hdul[0].header['METHOD'] - containers.camera = hdul[0].header['CAMERA'] - - containers.pixels_id = hdul[0].data - containers.charges_hg = hdul[1].data - containers.charges_lg = hdul[2].data - containers.peak_hg = hdul[3].data - containers.peak_lg = hdul[4].data - + log.info(f"Loading {filename}") + else: + log.info(f"Loading in {path} run number {run_number}") + filename = Path(path) / f"charge_run{run_number}.fits" + + with fits.open(filename) as hdul: + container = ChargesContainer() + container.run_number = hdul[0].header['RUN'] + container.nevents = hdul[0].header['NEVENTS'] + container.npixels = hdul[0].header['NPIXELS'] + container.method = hdul[0].header['METHOD'] + container.camera = hdul[0].header['CAMERA'] + container.pixels_id = hdul[0].data + container.charges_hg = hdul[1].data + container.charges_lg = hdul[2].data + container.peak_hg = hdul[3].data + container.peak_lg = hdul[4].data broken_pixels = hdul[5].data - containers.broken_pixels_hg = broken_pixels["HG broken pixels"] - containers.broken_pixels_lg = broken_pixels["LG broken pixels"] - - - + container.broken_pixels_hg = broken_pixels["HG broken pixels"] + container.broken_pixels_lg = broken_pixels["LG broken pixels"] table_prop = hdul[6].data - containers.event_id = table_prop["event_id"] - containers.event_type = table_prop["event_type"] - containers.ucts_timestamp = table_prop["ucts_timestamp"] - containers.ucts_busy_counter = table_prop["ucts_busy_counter"] - containers.ucts_event_counter = table_prop["ucts_event_counter"] - containers.multiplicity = table_prop["multiplicity"] - + container.event_id = table_prop["event_id"] + container.event_type = table_prop["event_type"] + container.ucts_timestamp = table_prop["ucts_timestamp"] + container.ucts_busy_counter = table_prop["ucts_busy_counter"] + container.ucts_event_counter = table_prop["ucts_event_counter"] + container.multiplicity = table_prop["multiplicity"] table_trigger = hdul[7].data - containers.trig_pattern_all = table_trigger["trig_pattern_all"] - containers.trig_pattern = table_trigger["trig_pattern"] - return containers + container.trig_pattern_all = table_trigger["trig_pattern_all"] + container.trig_pattern = table_trigger["trig_pattern"] + return container diff --git a/src/nectarchain/data/container/core.py b/src/nectarchain/data/container/core.py index 6471b3a6..d54140c4 100644 --- a/src/nectarchain/data/container/core.py +++ b/src/nectarchain/data/container/core.py @@ -7,6 +7,7 @@ from ctapipe.containers import Container,Field import numpy as np +__all__= ["ArrayDataContainer"] class ArrayDataContainer(Container): """ diff --git a/src/nectarchain/data/container/tests/test_charge.py b/src/nectarchain/data/container/tests/test_charge.py index 8018ed09..b15ebaa1 100644 --- a/src/nectarchain/data/container/tests/test_charge.py +++ b/src/nectarchain/data/container/tests/test_charge.py @@ -1,43 +1,51 @@ -from nectarchain.data.container.charge import ChargeContainers, ChargeContainer -from nectarchain.data.container.waveforms import WaveformsContainers +from ..chargesContainer import ChargesContainer,ChargesContainerIO +from ..waveformsContainer import WaveformsContainer +from ....makers import ChargesMaker import glob import numpy as np def create_fake_chargeContainer() : - pixels_id = np.array([2,4,3,8,6,9,7,1,5,10]) - nevents = 40 - npixels = 10 - rng = np.random.default_rng() - charge_hg = rng.integers(low=0, high=1000, size= (nevents,npixels)) - charge_lg = rng.integers(low=0, high=1000, size= (nevents,npixels)) - peak_hg = rng.integers(low=0, high=60, size= (nevents,npixels)) - peak_lg = rng.integers(low=0, high=60, size= (nevents,npixels)) - run_number = 1234 - return ChargeContainer( - charge_hg = charge_hg , - charge_lg = charge_lg, - peak_hg = peak_hg, - peak_lg = peak_lg, - run_number = run_number, - pixels_id = pixels_id, - nevents = nevents, - npixels = npixels + nevents = TestChargesContainer.nevents, + npixels = TestChargesContainer.npixels, + rng = np.random.default_rng(), + return ChargesContainer( + pixels_id = np.array([2,4,3,8,6,9,7,1,5,10]), + nevents =nevents, + npixels =npixels, + charge_hg = rng.integers(low=0, high=1000, size= (nevents,npixels)), + charge_lg = rng.integers(low=0, high=1000, size= (nevents,npixels)), + peak_hg = rng.integers(low=0, high=60, size= (nevents,npixels)), + peak_lg = rng.integers(low=0, high=60, size= (nevents,npixels)), + run_number = TestChargesContainer.run_number, + camera = 'TEST', + broken_pixels_hg = rng.integers(low=0, high=1, size= (nevents,npixels)), + broken_pixels_lg = rng.integers(low=0, high=1, size= (nevents,npixels)), + ucts_timestamp =rng.integers(low=0, high=100, size= (nevents)), + ucts_busy_counter =rng.integers(low=0, high=100, size= (nevents)), + ucts_event_counter =rng.integers(low=0, high=100, size= (nevents)), + event_type =rng.integers(low=0, high=1, size= (nevents)), + event_id =rng.integers(low=0, high=1000, size= (nevents)), + trig_pattern_all = rng.integers(low=0, high=1, size= (nevents,npixels,4)), + trig_pattern = rng.integers(low=0, high=1, size= (nevents,npixels)), + multiplicity =rng.integers(low=0, high=1, size= (nevents)) ) -class TestChargeContainer: - +class TestChargesContainer: + run_number = 1234 + nevents = 140 + npixels = 10 # Tests that a ChargeContainer object can be created with valid input parameters. def test_create_charge_container(self): pixels_id = np.array([2,4,3,8,6,9,7,1,5,10]) - nevents = 40 - npixels = 10 + nevents = TestChargesContainer.nevents + npixels = TestChargesContainer.npixels + run_number = TestChargesContainer.run_number charge_hg = np.random.randn(nevents,npixels) charge_lg = np.random.randn(nevents,npixels) peak_hg = np.random.randn(nevents,npixels) peak_lg = np.random.randn(nevents,npixels) - run_number = 1234 method = 'FullWaveformSum' - charge_container = ChargeContainer( + charge_container = ChargesContainer( charge_hg = charge_hg , charge_lg = charge_lg, peak_hg = peak_hg, @@ -49,8 +57,8 @@ def test_create_charge_container(self): method = method ) - assert np.allclose(charge_container.charge_hg,charge_hg) - assert np.allclose(charge_container.charge_lg,charge_lg) + assert np.allclose(charge_container.charges_hg,charge_hg) + assert np.allclose(charge_container.charges_lg,charge_lg) assert np.allclose(charge_container.peak_hg,peak_hg) assert np.allclose(charge_container.peak_lg,peak_lg) assert charge_container.run_number == run_number @@ -71,56 +79,39 @@ def test_create_charge_container(self): # Tests that the ChargeContainer object can be written to a file and the file is created. def test_write_charge_container(self, tmp_path = "/tmp"): charge_container = create_fake_chargeContainer() + tmp_path += f"/{np.random.randn(1)}" - charge_container.write(tmp_path) + ChargesContainerIO.write(tmp_path,charge_container) - assert len(glob.glob(f"{tmp_path}/charge_run1234.fits")) == 1 + assert len(glob.glob(f"{tmp_path}/charge_run{TestChargesContainer.run_number}.fits")) == 1 # Tests that a ChargeContainer object can be loaded from a file and the object is correctly initialized. def test_load_charge_container(self, tmp_path = "/tmp"): - pixels_id = np.array([2,4,3,8,6,9,7,1,5,10]) - nevents = 40 - npixels = 10 - charge_hg = np.random.randn(nevents,npixels) - charge_lg = np.random.randn(nevents,npixels) - peak_hg = np.random.randn(nevents,npixels) - peak_lg = np.random.randn(nevents,npixels) - run_number = np.random.randn(1)[0] - method = 'FullWaveformSum' - charge_container = ChargeContainer( - charge_hg = charge_hg , - charge_lg = charge_lg, - peak_hg = peak_hg, - peak_lg = peak_lg, - run_number = run_number, - pixels_id = pixels_id, - nevents = nevents, - npixels = npixels, - method = method - ) + charge_container = create_fake_chargeContainer() + tmp_path += f"/{np.random.randn(1)}" - charge_container.write(tmp_path) + ChargesContainerIO.write(tmp_path,charge_container) - loaded_charge_container = ChargeContainer.from_file(tmp_path, run_number) + loaded_charge_container = ChargesContainerIO.load(tmp_path,TestChargesContainer.run_number ) - assert isinstance(loaded_charge_container, ChargeContainer) - assert np.allclose(loaded_charge_container.charge_hg,charge_hg) - assert np.allclose(loaded_charge_container.charge_lg,charge_lg) - assert np.allclose(loaded_charge_container.peak_hg,peak_hg) - assert np.allclose(loaded_charge_container.peak_lg,peak_lg) - assert loaded_charge_container.run_number == run_number - assert loaded_charge_container.pixels_id.tolist() == pixels_id.tolist() - assert loaded_charge_container.nevents == nevents - assert loaded_charge_container.npixels == npixels - assert loaded_charge_container.method == method + assert isinstance(loaded_charge_container, ChargesContainer) + assert np.allclose(loaded_charge_container.charges_hg,charge_container.charges_hg) + assert np.allclose(loaded_charge_container.charges_lg,charge_container.charges_lg) + assert np.allclose(loaded_charge_container.peak_hg,charge_container.peak_hg) + assert np.allclose(loaded_charge_container.peak_lg,charge_container.peak_lg) + assert loaded_charge_container.run_number == charge_container.run_number + assert loaded_charge_container.pixels_id.tolist() == charge_container.pixels_id.tolist() + assert loaded_charge_container.nevents == charge_container.nevents + assert loaded_charge_container.npixels == charge_container.npixels + assert loaded_charge_container.method == charge_container.method # Tests that the ChargeContainer object can be sorted by event_id and the object is sorted accordingly. def test_sort_charge_container(self): charge_container = create_fake_chargeContainer() - charge_container.sort() + sorted_charge_container = ChargesMaker.sort(charge_container) - assert charge_container.event_id.tolist() == sorted(charge_container.event_id.tolist()) + assert sorted_charge_container.event_id.tolist() == sorted(charge_container.event_id.tolist()) # Tests that the run_number, pixels_id, npixels, nevents, method, multiplicity, and trig_pattern properties of the ChargeContainer object can be accessed and the values are correct. def test_access_properties(self): @@ -133,7 +124,7 @@ def test_access_properties(self): peak_lg = np.random.randn(nevents,npixels) run_number = 1234 method = 'FullWaveformSum' - charge_container = ChargeContainer( + charge_container = ChargesContainer( charge_hg = charge_hg , charge_lg = charge_lg, peak_hg = peak_hg, @@ -151,133 +142,4 @@ def test_access_properties(self): assert charge_container.nevents == nevents assert charge_container.method == method assert charge_container.multiplicity.shape == (nevents,) - assert charge_container.trig_pattern.shape == (nevents,4) - - -class TestChargeContainers: - - # Tests that an instance of ChargeContainers can be created with default arguments - def test_create_instance_with_default_arguments(self): - charge_containers = ChargeContainers() - assert len(charge_containers.chargeContainers) == 0 - assert charge_containers.nChargeContainer == 0 - - # Tests that an instance of ChargeContainers can be created from a WaveformsContainers instance - #def test_create_instance_from_waveforms(self): - # waveform_containers = WaveformsContainers(run_number = 1234) - # charge_containers = ChargeContainers.from_waveforms(waveform_containers) - # assert len(charge_containers.chargeContainers) == waveform_containers.nWaveformsContainer - # assert charge_containers.nChargeContainer == waveform_containers.nWaveformsContainer - - # Tests that ChargeContainers can be written to disk - def test_write_to_disk(self, tmpdir): - charge_containers = ChargeContainers() - - charge_container = create_fake_chargeContainer() - - charge_containers.append(charge_container) - path = str(tmpdir) - charge_containers.write(path) - assert len(glob.glob(f"{path}/charge_run*1234*.fits")) == 1 - - # Tests that ChargeContainers can be loaded from disk - def test_load_from_disk(self, tmpdir): - charge_containers = ChargeContainers() - - pixels_id = np.array([2,4,3,8,6,9,7,1,5,10]) - nevents = 40 - npixels = 10 - charge_hg = np.random.randn(nevents,npixels) - charge_lg = np.random.randn(nevents,npixels) - peak_hg = np.random.randn(nevents,npixels) - peak_lg = np.random.randn(nevents,npixels) - run_number = np.random.randn(1)[0] - - charge_container = ChargeContainer( - charge_hg = charge_hg , - charge_lg = charge_lg, - peak_hg = peak_hg, - peak_lg = peak_lg, - run_number = run_number, - pixels_id = pixels_id, - nevents = nevents, - npixels = npixels - ) - charge_containers.append(charge_container) - path = str(tmpdir) - charge_containers.write(path) - loaded_charge_containers = ChargeContainers.from_file(path, run_number) - assert len(loaded_charge_containers.chargeContainers) == 1 - assert loaded_charge_containers.nChargeContainer == 1 - - # Tests that a ChargeContainer can be appended to a ChargeContainers instance - def test_append_charge_container(self): - charge_containers = ChargeContainers() - pixels_id = np.array([2,4,3,8,6,9,7,1,5,10]) - nevents = 40 - npixels = 10 - charge_hg = np.random.randn(nevents,npixels) - charge_lg = np.random.randn(nevents,npixels) - peak_hg = np.random.randn(nevents,npixels) - peak_lg = np.random.randn(nevents,npixels) - run_number = 1234 - - charge_container = ChargeContainer( - charge_hg = charge_hg , - charge_lg = charge_lg, - peak_hg = peak_hg, - peak_lg = peak_lg, - run_number = run_number, - pixels_id = pixels_id, - nevents = nevents, - npixels = npixels - ) - charge_containers.append(charge_container) - assert len(charge_containers.chargeContainers) == 1 - assert charge_containers.nChargeContainer == 1 - - # Tests that a ChargeContainers instance can be merged into a single ChargeContainer - def test_merge_into_single_charge_container(self): - charge_containers = ChargeContainers() - pixels_id = np.array([2,4,3,8,6,9,7,1,5,10]) - nevents = 40 - npixels = 10 - charge_hg = np.random.randn(nevents,npixels) - charge_lg = np.random.randn(nevents,npixels) - peak_hg = np.random.randn(nevents,npixels) - peak_lg = np.random.randn(nevents,npixels) - run_number = 1234 - - charge_container1 = ChargeContainer( - charge_hg = charge_hg , - charge_lg = charge_lg, - peak_hg = peak_hg, - peak_lg = peak_lg, - run_number = run_number, - pixels_id = pixels_id, - nevents = nevents, - npixels = npixels - ) - pixels_id = np.array([2,4,3,8,6,9,7,1,5,10]) - nevents = 40 - npixels = 10 - charge_hg = np.random.randn(nevents,npixels) - charge_lg = np.random.randn(nevents,npixels) - peak_hg = np.random.randn(nevents,npixels) - peak_lg = np.random.randn(nevents,npixels) - run_number = 1234 - - charge_container2 = ChargeContainer( - charge_hg = charge_hg , - charge_lg = charge_lg, - peak_hg = peak_hg, - peak_lg = peak_lg, - run_number = run_number, - pixels_id = pixels_id, - nevents = nevents, - npixels = npixels - ) - charge_containers.append(charge_container1) - charge_containers.append(charge_container2) - merged_charge_container = charge_containers.merge() - assert merged_charge_container.nevents == charge_container1.nevents + charge_container2.nevents \ No newline at end of file + assert charge_container.trig_pattern.shape == (nevents,4) \ No newline at end of file diff --git a/src/nectarchain/data/container/tests/test_waveforms.py b/src/nectarchain/data/container/tests/test_waveforms.py index e69de29b..d3689f79 100644 --- a/src/nectarchain/data/container/tests/test_waveforms.py +++ b/src/nectarchain/data/container/tests/test_waveforms.py @@ -0,0 +1,76 @@ +from ..waveformsContainer import WaveformsContainer,WaveformsContainerIO +from ....makers import WaveformsMaker +import glob +import numpy as np + +def create_fake_waveformsContainer() : + nevents = TestWaveformsContainer.nevents, + npixels = TestWaveformsContainer.npixels, + nsamples = TestWaveformsContainer.nsamples, + rng = np.random.default_rng(), + + return WaveformsContainer( + pixels_id = np.array([2,4,3,8,6,9,7,1,5,10]), + nevents =nevents, + npixels =npixels, + wfs_hg = rng.integers(low=0, high=1000, size= (nevents,npixels,nsamples)), + wfs_lg = rng.integers(low=0, high=1000, size= (nevents,npixels,nsamples)), + run_number = TestWaveformsContainer.run_number, + camera = 'TEST', + broken_pixels_hg = rng.integers(low=0, high=1, size= (nevents,npixels)), + broken_pixels_lg = rng.integers(low=0, high=1, size= (nevents,npixels)), + ucts_timestamp =rng.integers(low=0, high=100, size= (nevents)), + ucts_busy_counter =rng.integers(low=0, high=100, size= (nevents)), + ucts_event_counter =rng.integers(low=0, high=100, size= (nevents)), + event_type =rng.integers(low=0, high=1, size= (nevents)), + event_id =rng.integers(low=0, high=1000, size= (nevents)), + trig_pattern_all = rng.integers(low=0, high=1, size= (nevents,npixels,4)), + trig_pattern = rng.integers(low=0, high=1, size= (nevents,npixels)), + multiplicity =rng.integers(low=0, high=1, size= (nevents)) + ) + +class TestWaveformsContainer: + run_number = 1234 + nevents = 140 + npixels = 10 + nsamples = 5 + # Tests that a ChargeContainer object can be created with valid input parameters. + def test_create_waveform_container(self): + waveform_container = create_fake_waveformsContainer() + assert isinstance(waveform_container,WaveformsContainer) + + + # Tests that the ChargeContainer object can be written to a file and the file is created. + def test_write_waveform_container(self, tmp_path = "/tmp"): + waveform_container = create_fake_waveformsContainer() + tmp_path += f"/{np.random.randn(1)}" + + WaveformsContainerIO.write(tmp_path,waveform_container) + + assert len(glob.glob(f"{tmp_path}/*_run{TestWaveformsContainer.run_number}.fits")) == 1 + + # Tests that a ChargeContainer object can be loaded from a file and the object is correctly initialized. + def test_load_waveform_container(self, tmp_path = "/tmp"): + waveform_container = create_fake_waveformsContainer() + tmp_path += f"/{np.random.randn(1)}" + + WaveformsContainerIO.write(tmp_path,waveform_container) + + loaded_waveform_container = WaveformsContainerIO.load(tmp_path,TestWaveformsContainer.run_number ) + + assert isinstance(loaded_waveform_container, WaveformsContainer) + assert np.allclose(loaded_waveform_container.wfs_hg,waveform_container.charges_hg) + assert np.allclose(loaded_waveform_container.wfs_lg,waveform_container.charges_lg) + assert loaded_waveform_container.run_number == waveform_container.run_number + assert loaded_waveform_container.pixels_id.tolist() == waveform_container.pixels_id.tolist() + assert loaded_waveform_container.nevents == waveform_container.nevents + assert loaded_waveform_container.npixels == waveform_container.npixels + assert loaded_waveform_container.nsamples == waveform_container.nsamples + + # Tests that the ChargeContainer object can be sorted by event_id and the object is sorted accordingly. + def test_sort_waveform_container(self): + waveform_container = create_fake_waveformsContainer() + + sorted_waveform_container = WaveformsMaker.sort(waveform_container) + + assert sorted_waveform_container.event_id.tolist() == sorted(waveform_container.event_id.tolist()) \ No newline at end of file diff --git a/src/nectarchain/data/container/waveformsContainer.py b/src/nectarchain/data/container/waveformsContainer.py index 1a277b11..57563e80 100644 --- a/src/nectarchain/data/container/waveformsContainer.py +++ b/src/nectarchain/data/container/waveformsContainer.py @@ -26,7 +26,7 @@ class WaveformsContainer(ArrayDataContainer): """ A container that holds information about waveforms from a specific run. - Attributes: + Fields: nsamples (int): The number of samples in the waveforms. subarray (SubarrayDescription): The subarray description instance. wfs_hg (np.ndarray): An array of high gain waveforms. @@ -51,15 +51,55 @@ class WaveformsContainer(ArrayDataContainer): ) + class WaveformsContainerIO(ABC) : + """ + The `WaveformsContainerIO` class provides methods for writing and loading `WaveformsContainer` instances to/from FITS files. It also includes a method for writing the subarray configuration to an HDF5 file. + + Example Usage: + # Writing a WaveformsContainer instance to a FITS file + container = WaveformsContainer() + # ... populate the container with data ... + WaveformsContainerIO.write('/path/to/output', container) + + # Loading a WaveformsContainer instance from a FITS file + container = WaveformsContainerIO.load('/path/to/input.fits') + + Main functionalities: + - Writing a `WaveformsContainer` instance to a FITS file, including the waveforms data and metadata. + - Loading a `WaveformsContainer` instance from a FITS file, including the waveforms data and metadata. + - Writing the subarray configuration to an HDF5 file. + + Methods: + - `write(path: str, containers: WaveformsContainer, **kwargs) -> None`: Writes a `WaveformsContainer` instance to a FITS file. The method also creates an HDF5 file representing the subarray configuration. + - `load(path: str) -> WaveformsContainer`: Loads a `WaveformsContainer` instance from a FITS file. The method also loads the subarray configuration from the corresponding HDF5 file. + + Fields: + None. + """ @staticmethod - def write(path : str, containers : WaveformsContainer, **kwargs): - '''method to write in an output FITS file the WaveformsContainer. Two files are created, one FITS representing the data - and one HDF5 file representing the subarray configuration + def write(path : str, containers : WaveformsContainer, **kwargs) -> None: + '''Write the WaveformsContainer data to an output FITS file. + + This method creates two files: one FITS file representing the waveform data and one HDF5 file representing the subarray configuration. Args: - path (str): the directory where you want to save data + path (str): The directory where you want to save the data. + containers (WaveformsContainer): The WaveformsContainer instance containing the data to be saved. + **kwargs: Additional keyword arguments for customization (optional). + + Keyword Args: + suffix (str, optional): A suffix to append to the output file names (default is ''). + overwrite (bool, optional): If True, overwrite the output files if they already exist (default is False). + Returns: + None: This method does not return any value. + Raises: + OSError: If there is an error while writing the FITS file. + Exception: If there is any other exception during the writing process. + Example: + waveformsContainer = WaveformsContainer() + WaveformsContainerIO.write(path, waveformsContainer, suffix="v1", overwrite=True) ''' suffix = kwargs.get("suffix","") if suffix != "" : suffix = f"_{suffix}" @@ -118,16 +158,20 @@ def write(path : str, containers : WaveformsContainer, **kwargs): raise e @staticmethod - def load(path : str) : - '''load WaveformsContainer from FITS file previously written with WaveformsContainer.write() method - Note : 2 files are loaded, the FITS one representing the waveforms data and a HDF5 file representing the subarray configuration. - This second file has to be next to the FITS file. - + def load(path : str) -> WaveformsContainer: + '''Load a WaveformsContainer from a FITS file previously written with WaveformsContainerIO.write() method. + + Note: Two files are loaded—the FITS file representing the waveform data and an HDF5 file representing the subarray configuration. + The HDF5 file should be located next to the FITS file. + Args: - path (str): path of the FITS file + path (str): The path to the FITS file containing the waveform data. Returns: - WaveformsContainer: WaveformsContainer instance + WaveformsContainer: A WaveformsContainer instance loaded from the specified file. + Example: + waveformsContainer = WaveformsContainerIO.load(path, run_number) + """ ''' log.info(f"loading from {path}") with fits.open(Path(path)) as hdul :