diff --git a/nenupy/__init__.py b/nenupy/__init__.py index 4eb3110..1beb82c 100644 --- a/nenupy/__init__.py +++ b/nenupy/__init__.py @@ -5,7 +5,7 @@ __copyright__ = "Copyright 2023, nenupy" __credits__ = ["Alan Loh"] __license__ = "MIT" -__version__ = "2.6.0" +__version__ = "2.6.1" __maintainer__ = "Alan Loh" __email__ = "alan.loh@obspm.fr" diff --git a/nenupy/beamlet/sdata.py b/nenupy/beamlet/sdata.py index 3ae1630..9072c1f 100644 --- a/nenupy/beamlet/sdata.py +++ b/nenupy/beamlet/sdata.py @@ -546,11 +546,12 @@ def fbackground(self): # --------------------------------------------------------- # # ------------------------ Methods ------------------------ # - def plot(self, polarization=None, figname=None, db=True, **kwargs): + def plot(self, fig=None, ax=None, polarization=None, figname=None, db=True, **kwargs): """ kwargs keys: cmap, title, cblabel, figsize, altaza, vmin, vmax """ import matplotlib.pyplot as plt + import matplotlib.dates as mdates if polarization is None: pol_idx = 0 @@ -573,10 +574,17 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs): if 'figsize' not in kwargs.keys(): kwargs['figsize'] = (15, 10) - fig = plt.figure(figsize=kwargs['figsize']) + if ax is None: + # Create a full figure from the start + fig = plt.figure(figsize=kwargs['figsize']) + ax = fig.add_subplot() + else: + # Fill up the input ax with the plot + pass + if len(dynspec.shape) == 1: if dynspec.size == self.datetime.size: - plt.plot( + ax.plot( self.datetime, dynspec ) @@ -585,20 +593,20 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs): for ptime in ptimes: if (ptime < self.datetime[0]) or (ptime > self.datetime[-1]): continue - plt.axvline(ptime.datetime, linestyle='-.', color='black') - plt.ylim((kwargs.get('vmin', None), kwargs.get('vmax', None))) - plt.xlabel( - f'Time (since {self.time[0].isot})' + ax.axvline(ptime.datetime, linestyle='-.', color='black') + ax.ylim((kwargs.get('vmin', None), kwargs.get('vmax', None))) + ax.set_xlabel( + f'Time (UTC from {self.time[0].isot})' ) - plt.ylabel(kwargs['cblabel']) + ax.set_ylabel(kwargs['cblabel']) elif dynspec.size == self.freq.size: - plt.plot( + ax.plot( self.freq.to(u.MHz).value, dynspec ) - plt.ylim((kwargs.get('vmin', None), kwargs.get('vmax', None))) - plt.xlabel('Frequency (MHz)') - plt.ylabel(kwargs['cblabel']) + ax.set_ylim((kwargs.get('vmin', None), kwargs.get('vmax', None))) + ax.set_xlabel('Frequency (MHz)') + ax.set_ylabel(kwargs['cblabel']) else: if 'vmin' not in kwargs.keys(): kwargs['vmin'] = np.nanpercentile(dynspec, 5) @@ -612,7 +620,7 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs): kwargs['vmax'] = np.nanpercentile(dynspec, 95) else: pass - plt.pcolormesh( + im = ax.pcolormesh( self.datetime, self.freq.to(u.MHz).value, dynspec, @@ -626,16 +634,16 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs): for ptime in ptimes: if (ptime < self.datetime[0]) or (ptime > self.datetime[-1]): continue - plt.axvline(ptime.datetime, linestyle='-.', color='black') - cbar = plt.colorbar(pad=0.03)#format='%.1e') + ax.axvline(ptime.datetime, linestyle='-.', color='black') + cbar = plt.colorbar(im, pad=0.03)#format='%.1e') cbar.set_label(kwargs['cblabel']) if kwargs.get("overlay", None) is not None: - ax = plt.gca() + #ax = plt.gca() xlim = ax.get_xlim() ylim = ax.get_ylim() overlay_time, overlay_frequency, overlay_values = kwargs["overlay"] - plt.pcolor( + ax.pcolor( overlay_time.datetime, overlay_frequency.to(u.MHz).value, overlay_values, @@ -655,17 +663,30 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs): ax.set_xlim(xlim) ax.set_ylim(ylim) - plt.xlabel( - f'Time (since {self.time[0].isot})' + ax.set_xlabel( + f'Time (UTC from {self.time[0].isot})' ) - plt.ylabel('Frequency (MHz)') - plt.title(kwargs['title']) - + ax.set_ylabel('Frequency (MHz)') + ax.set_title(kwargs['title']) + + # Set x axis labels + # ax.xaxis.set_minor_locator(mdates.MinuteLocator(byminute=[15, 30, 45])) + # ax.xaxis.set_major_locator(mdates.HourLocator()) + # hourFmt = mdates.DateFormatter("%H", usetex=True) + # ax.xaxis.set_major_formatter(hourFmt) + locator = ax.xaxis.set_major_locator(mdates.AutoDateLocator()) + ax.xaxis.set_major_formatter( + mdates.ConciseDateFormatter(locator, show_offset=False) + ) + # Save or show + if not (ax is None): + return + if figname is None: plt.show() - elif figname.lower() == 'return': - return fig + elif (figname is None) or (figname == ""): + return else: fig.savefig( figname, @@ -674,7 +695,6 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs): bbox_inches='tight' ) plt.close('all') - return # --------------------------------------------------------- # diff --git a/nenupy/io/tf.py b/nenupy/io/tf.py index ab06fc6..b420870 100644 --- a/nenupy/io/tf.py +++ b/nenupy/io/tf.py @@ -113,8 +113,8 @@ def func_call(self) -> Callable: @classmethod def correct_bandpass(cls): - """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.correct_bandpass` to correct the polyphase-filter bandpass reponse. - """ + """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.correct_bandpass` to correct the polyphase-filter bandpass reponse.""" + def wrapper_task(data, channels): return utils.correct_bandpass(data=data, n_channels=channels) @@ -122,8 +122,8 @@ def wrapper_task(data, channels): @classmethod def remove_channels(cls): - """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.remove_channels_per_subband` to set a list of sub-band channels to `NaN` values. - """ + """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.remove_channels_per_subband` to set a list of sub-band channels to `NaN` values.""" + def wrapper_task(data, channels, remove_channels): if (remove_channels is None) or (len(remove_channels) == 0): return data @@ -180,8 +180,8 @@ def wrapper_task( @classmethod def correct_faraday_rotation(cls): - """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.de_faraday_data` to correct for Faraday rotation for a given ``'rotation_measure'`` set in :attr:`~nenupy.io.tf.TFPipeline.parameters`. - """ + """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.de_faraday_data` to correct for Faraday rotation for a given ``'rotation_measure'`` set in :attr:`~nenupy.io.tf.TFPipeline.parameters`.""" + def apply_faraday(frequency_hz, data, rotation_measure): if rotation_measure is None: return frequency_hz, data @@ -197,17 +197,20 @@ def apply_faraday(frequency_hz, data, rotation_measure): def de_disperse(cls): """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.de_disperse_array` to de-disperse the data using the ``'dispersion_measure'`` set in :attr:`~nenupy.io.tf.TFPipeline.parameters`. - .. warning:: + .. warning:: - Due to the configuration of the underlying :class:`~dask.array.core.Array`, its :meth:`dask.array.Array.compute` method has to be applied priori to de-dispersing the data. - Therefore, a potential huge data volume may be computed at once. - By default, a security exception is raised to prevent computing a too large data set. - To bypass this limit, set ``'ignore_volume_warning'`` of :attr:`~nenupy.io.tf.TFPipeline.parameters` to `True`. + Due to the configuration of the underlying :class:`~dask.array.core.Array`, its :meth:`dask.array.Array.compute` method has to be applied priori to de-dispersing the data. + Therefore, a potential huge data volume may be computed at once. + By default, a security exception is raised to prevent computing a too large data set. + To bypass this limit, set ``'ignore_volume_warning'`` of :attr:`~nenupy.io.tf.TFPipeline.parameters` to `True`. """ - def wrapper_task(frequency_hz, data, dt, dispersion_measure, ignore_volume_warning): + + def wrapper_task( + frequency_hz, data, dt, dispersion_measure, ignore_volume_warning + ): if dispersion_measure is None: - return frequency_hz, data + return frequency_hz, data # Make sure the data volume is not too big! projected_data_volume = data.nbytes * u.byte if (projected_data_volume >= DATA_VOLUME_SECURITY_THRESHOLD) and ( @@ -228,7 +231,11 @@ def wrapper_task(frequency_hz, data, dt, dispersion_measure, ignore_volume_warni ) return frequency_hz, da.from_array(data, chunks=tmp_chuncks) - return cls("De-disperse", wrapper_task, ["dt", "dispersion_measure", "ignore_volume_warning"]) + return cls( + "De-disperse", + wrapper_task, + ["dt", "dispersion_measure", "ignore_volume_warning"], + ) @classmethod def time_rebin(cls): @@ -338,7 +345,7 @@ def __init__(self, data_obj: Any, *tasks: TFTask): def __repr__(self) -> str: return self.info() - + @property def parameters(self) -> utils.TFPipelineParameters: """_summary_ @@ -347,6 +354,7 @@ def parameters(self) -> utils.TFPipelineParameters: :rtype: :class:`~nenupy.io.tf_utils.TFPipelineParameters` """ return self._parameters + @parameters.setter def parameters(self, params: utils.TFPipelineParameters) -> None: self._parameters = params @@ -545,7 +553,9 @@ def info(self) -> None: print(message) def get(self, **pipeline_kwargs) -> SData: - """_summary_ + """Select time-frequency data, run the user-defined pipeline and return the product. + Data selection, as well as pipeline specific arguments are defined as keyword arguments and passed to :attr:`nenupy.io.tf.TFPipeline.parameters`. + The pipeline can be accessed and modified through the :attr:`nenupy.io.tf.Spectra.pipeline` attribute. .. rubric:: Available parameters @@ -553,37 +563,37 @@ def get(self, **pipeline_kwargs) -> SData: :param tmin: Lower edge of time selection, can either be given as a :class:`~astropy.time.Time` object or an ISOT/ISO string. :type tmin: `str` or :class:`~astropy.time.Time` - :param tmax: Hello + :param tmax: Upper edge of time selection, can either be given as an :class:`~astropy.time.Time` object or an ISOT/ISO string. :type tmax: `str` or :class:`~astropy.time.Time` - :param fmin: Hello - :type fmin: `str` or :class:`~astropy.unit.Quantity` - :param fmin: Hello - :type fmax: `str` or :class:`~astropy.unit.Quantity` - :param beam: Hello - :type beam: str or :class:`~astropy.time.Time` - :param dispersion_measure: Hello + :param fmin: Lower frequency boundary selection, can either be given as a :class:`~astropy.unit.Quantity` object or float (assumed to be in MHz in that case). + :type fmin: `float` or :class:`~astropy.unit.Quantity` + :param fmax: Higher frequency boundary selection, can either be given as a :class:`~astropy.unit.Quantity` object or float (assumed to be in MHz in that case). + :type fmax: `float` or :class:`~astropy.unit.Quantity` + :param beam: Beam selection, a single integer corresponding to the index of a recorded numerical beam is expected. Default is the first recorded. + :type beam: `int` + :param dispersion_measure: Enable de-dispersion of the data by this Dispersion Measure. Note that the :meth:`~nenupy.io.tf.TFTask.de_disperse` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`). It can either be provided as a :class:`~astropy.Quantity` object or a float (assumed to be in pc/cm^3 in that case). :type dispersion_measure: `float` or :class:`~astropy.unit.Quantity` :param rotation_measure: Hello :type rotation_measure: `float` or :class:`~astropy.unit.Quantity` - :param rebin_dt: Hello + :param rebin_dt: Desired rebinning time resolution, can either be given as a :class:`~astropy.unit.Quantity` object or a float (assumed to be in sec in that case). Note that the :meth:`~nenupy.io.tf.TFTask.time_rebin` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`). :type rebin_dt: `float` or :class:`~astropy.unit.Quantity` - :param rebin_df: Hello + :param rebin_df: Desired rebinning frequency resolution, can either be given as a :class:`~astropy.unit.Quantity` object or float (assumed to be in kHz in that case). Note that the :meth:`~nenupy.io.tf.TFTask.frequency_rebin` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`). :type rebin_df: `float` or :class:`~astropy.unit.Quantity` - :param remove_channels: Hello - :type remove_channels: str or :class:`~astropy.time.Time` - :param dreambeam_skycoord: Hello - :type dreambeam_skycoord: str or :class:`~astropy.time.Time` - :param dreambeam_dt: Hello - :type dreambeam_dt: str or :class:`~astropy.time.Time` - :param dreambeam_parallactic: Hello - :type dreambeam_parallactic: bool - :param stokes: Hello - :type stokes: str or :class:`~astropy.time.Time` - :param ignore_volume_warning: Hello - :type ignore_volume_warning: bool - - :return: _description_ - :rtype: SData + :param remove_channels: List of subband channels to remove, e.g. `remove_channels=[0,1,-1]` would remove the first, second (low-freq) and last channels from each subband. Note that the :meth:`~nenupy.io.tf.TFTask.remove_channels` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`). + :type remove_channels: `list` or :class:`~numpy.ndarray` + :param dreambeam_skycoord: Tracked celestial coordinates used during *DreamBeam* correction (along with ``'dreambeam_dt'`` and ``'dreambeam_parallactic'``), a :class:`~astropy.coordinates.SkyCoord` object is expected. Note that the :meth:`~nenupy.io.tf.TFTask.correct_polarization` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`). + :type dreambeam_skycoord: :class:`~astropy.coordinates.SkyCoord` + :param dreambeam_dt: *DreamBeam* correction time resolution (along with ``'dreambeam_skycoord'`` and ``'dreambeam_parallactic'``), a :class:`~astropy.Quantity` object or a float (assumed in seconds) are expected. Note that the :meth:`~nenupy.io.tf.TFTask.correct_polarization` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`). + :type dreambeam_dt: `float` or :class:`~astropy.unit.Quantity` + :param dreambeam_parallactic: *DreamBeam* parallactic angle correction (along with ``'dreambeam_skycoord'`` and ``'dreambeam_dt'``), a boolean is expected. Note that the :meth:`~nenupy.io.tf.TFTask.correct_polarization` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`). + :type dreambeam_parallactic: `bool` + :param stokes: Stokes parameter selection, can either be given as a string or a list of strings, e.g. ['I', 'Q', 'V/I']. Note that the :meth:`~nenupy.io.tf.TFTask.get_stokes` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`). + :type stokes: `str` or `list[str]` + :param ignore_volume_warning: Ignore or not (default value) the limit regarding output data volume. + :type ignore_volume_warning: `bool` + + :return: Processed data selection. + :rtype: :class:`~nenupy.beamlet.sdata.SData` """ # Update the pipeline parameters to user's last requests