Added fp_oversample

fp_tools/tools_3d.py

@@ -1,11 +1,157 @@
+import itertools
+import logging
+import multiprocessing
 import threading
 import sys
 
 from astropy.io import fits
 import numpy as np
 
 
-class z_repeat(threading.Thread):
+# noinspection PyUnusedLocal,PyUnusedLocal
+def signal_handler(s, frame):
+    sys.exit()
+
+
+class ZCut(threading.Thread):
+    """
+    Extract a part of the data-cube in the spectral direction and keeping the
+    whole spatial pixels. The header is updated to maintain the calibration
+    in the same direction.
+
+    Parameters
+    ----------
+        _input : str
+            The input cube filename.
+
+        _output : str
+            The output cube filename.
+
+        n_begin : int
+            The channel that will be considered as the first new one
+            (inclusive).
+
+        n_end : int
+            The channel that will be considered as the last new one
+            (exclusive).
+    """
+    def __init__(self, _input, _output, n_begin=None, n_end=None):
+
+        threading.Thread.__init__(self)
+
+        self._input = _input
+        self._output = _output
+        self.n_begin = n_begin
+        self.n_end = n_end
+
+    def run(self):
+
+        data = fits.getdata(self._input)
+        header = fits.getheader(self._input)
+
+        if self.n_begin is not None:
+            data = data[self.n_begin:]
+            header['CRPIX3'] -= (self.n_begin + 1)
+        else:
+            self.n_begin = 0
+
+        if self.n_end is not None:
+            data = data[:self.n_end]
+        else:
+            self.n_end = self.data.shape[0]
+
+        header.add_history('Cube cut from z = %i to z = %i' %
+                           (self.n_begin, self.n_end))
+
+        fits.writeto(self._output, data, header)
+
+
+class ZOversample(threading.Thread):
+    """
+    Oversample a data-cube in the spectral direction using linear fitting.
+
+    Parameters
+    ----------
+        _input : str
+            The input cube filename
+
+        _output : str
+            The output cube filename
+
+        oversample_factor : int
+            The oversample factor. In other words, how many points will be added
+            for each original point.
+    """
+    def __init__(self, _input, _output, oversample_factor):
+
+        threading.Thread.__init__(self)
+
+        self.input = _input
+        self.output = _output
+        self.oversample_factor = oversample_factor
+
+    def run(self):
+
+        # Read data
+        data = fits.getdata(self.input)
+        header = fits.getheader(self.input)
+
+        # Extract information
+        depth, height, width = data.shape
+
+        new_depth = self.oversample_factor * depth # - (self.oversample_factor - 1)
+        print(new_depth)
+
+        assert(isinstance(new_depth, int))
+        del data
+
+        x = np.arange(header['NAXIS1'], dtype=int)
+        y = np.arange(header['NAXIS2'], dtype=int)
+
+        # Keep the terminal alive
+        loading = multiprocessing.Process(target=stand_by, name="stand_by")
+        loading.start()
+
+        # Create a pool for subprocesses
+        p = multiprocessing.Pool(4)
+        results = None
+        fitter = OverSampler(self.input, new_depth)
+
+        try:
+            results = p.map_async(fitter, itertools.product(x, y)).get(99999999)
+        except KeyboardInterrupt:
+            print('\n\nYou pressed Ctrl+C!')
+            print('Leaving now. Bye!\n')
+            pass
+
+        # Closing processes
+        p.close()
+        p.join()
+        loading.terminate()
+
+        # Fix header
+        keys = ['CDELT3', 'C3_3', 'PHMSAMP']
+        for key in keys:
+            try:
+                header[key] /= self.oversample_factor
+            except KeyError:
+                print(' {:s} key is not present in the header.'.format(key))
+
+        try:
+            header['CRPIX3'] *= self.oversample_factor
+        except KeyError:
+            print(' {:s} key is not present in the header.'.format('CRPIX3'))
+
+        x = int(header['NAXIS1'])
+        y = int(header['NAXIS2'])
+        results = np.array(results)
+        results = results.reshape((x, y, new_depth))
+        results = np.transpose(results, axes=[2, 1, 0])
+
+        fits.writeto(self.output, results, header)
+
+
+class ZRepeat(threading.Thread):
     """
     Repeat a data-cube in the spectral direction (Z).
     """
@@ -59,34 +205,71 @@ def run(self):
         fits.writeto(self.output, temp_after, hdr)
 
 
-class z_cut(threading.Thread):
+class OverSampler:
 
-    def __init__(self, _input, _output, n_begin=None, n_end=None):
+    def __init__(self, filename, new_size):
+        """
+        Parameter
+        ---------
+            filename : str
+                Relative or absolute path to the file that contains a data-cube
+                from where the 2D maps will be extracted through gaussian
+                fitting.
+        """
+        self._filename = filename
+        self.new_size = new_size
 
-        threading.Thread.__init__(self)
+        # Load the data
+        self.data = fits.getdata(self._filename, memmap=True)
 
-        self._input = _input
-        self._output = _output
-        self.n_begin = n_begin
-        self.n_end = n_end
+    def __call__(self, indexes):
+        """
+        Parameter
+        ---------
+            indexes : tuple
+                Contains two integers that correspond to the X and Y indexes
+                that will be used to extract the spectrum from the data-cube and
+                fits a gaussian to this extracted spectrum.
+        Returns
+        -------
+            results : list
+                A list containing the numerical values of the three gaussian
+                parameters met in the fitting processes `peak`, `mean` and
+                `stddev`.
+        """
+        i, j = indexes
 
-    def run(self):
+        # Get the usefull information
+        s = self.data[:, j, i]
 
-        data = fits.getdata(self._input)
-        header = fits.getheader(self._input)
+        # Create new Z
+        z = np.arange(s.size)
+        z_interpolation = np.linspace(0, s.size, self.new_size)
 
-        if self.n_begin is not None:
-            data = data[self.n_begin:]
-            header['CRPIX3'] -= (self.n_begin + 1)
-        else:
-            self.n_begin = 0
+        # Interpolate
+        new_s = np.interp(z_interpolation, z, s)
 
-        if self.n_end is not None:
-            data = data[:self.n_end]
-        else:
-            self.n_end = self.data.shape[0]
+        return new_s
 
-        header.add_history('Cube cut from z = %i to z = %i' %
-                           (self.n_begin, self.n_end))
 
-        fits.writeto(self._output, data, header)
+def stand_by(level=logging.NOTSET):
+    """
+    A silly method that keeps the terminal alive so the user knows that
+    this programs is still running. :-)
+    """
+    from time import sleep
+
+    output = ['/', '-', '\\', '|']
+    i = 0
+
+    if level in [logging.NOTSET, logging.WARN, logging.ERROR]:
+        return
+
+    while True:
+        sys.stdout.write("\r [{:s}]".format(output[i]))
+        sys.stdout.flush()
+        sleep(0.5)
+        i += 1
+        i %= 4
+
+    return

fp_tools/version.py

@@ -1,9 +1,11 @@
 # Release notes
 # - Added history in the header.
+#
+# 0.3 - Added oversample
 
 api = 0
-feature = 2
-bug = 4
+feature = 3
+bug = 0
 
 __doc__ = '%d.%d.%d' % (api, feature, bug)
 

scripts/fp_cut

@@ -38,8 +38,8 @@ if os.path.exists(args.output_cube):
     sys.exit(1)
 
 # Running scripts
-fp_cut = fp_tools.tools_3d.z_cut(args.input_cube, args.output_cube,
-                                 n_begin=args.n_begin, n_end=args.n_end)
+fp_cut = fp_tools.tools_3d.ZCut(args.input_cube, args.output_cube,
+                                n_begin=args.n_begin, n_end=args.n_end)
 
 fp_cut.start()
 fp_cut.join()

scripts/fp_oversample

@@ -0,0 +1,39 @@
+#!python
+
+import os
+import sys
+import time
+
+import argparse
+import fp_tools
+
+
+# Parse command line arguments
+parser = argparse.ArgumentParser(
+    description="Oversample the data-cube in the Z direction using linear "
+                "interpolation.")
+
+parser.add_argument('input_cube', type=str, help="Input cube.")
+
+parser.add_argument('output_cube', type=str, help="Output cube.")
+
+parser.add_argument('oversample_factor', type=int, help="Oversample factor.")
+
+args = parser.parse_args()
+
+# Check before running
+if not os.path.exists(args.input_cube):
+    print("\n Input file does not exists: %s\n Leaving now." % args.input_cube)
+    sys.exit(1)
+
+if os.path.exists(args.output_cube):
+    print("\n Output file exists: %s" % args.output_cube)
+    print(" Delete it before running this.\n Leaving now.")
+    sys.exit(1)
+
+# Running scripts
+fp_repeat = fp_tools.tools_3d.ZOversample(args.input_cube, args.output_cube,
+                                          args.oversample_factor)
+
+fp_repeat.start()
+fp_repeat.join()

scripts/fp_repeat

@@ -37,8 +37,9 @@ if os.path.exists(args.output_cube):
     sys.exit(1)
 
 # Running scripts
-fp_repeat = fp_tools.tools_3d.z_repeat(args.input_cube, args.output_cube,
-                           n_after=args.end, n_before=args.begin)
+fp_repeat = fp_tools.tools_3d.ZRepeat(args.input_cube, args.output_cube,
+                                      n_after=args.end, n_before=args.begin)
 
 fp_repeat.start()
 fp_repeat.join()
+header['CDELT3'] /= self.oversample_factor

setup.py

@@ -41,7 +41,7 @@
 
     # package_dir={'': 'fp_tools'},
     # package_data={},
-    scripts=['scripts/fp_cut', 'scripts/fp_repeat'],
+    scripts=['scripts/fp_cut', 'scripts/fp_repeat', 'scripts/fp_oversample'],
     zip_safe=False,
 
     # Alternatively, if you want to distribute just a my_module.py, uncomment