source: sasview/src/sas/sascalc/file_converter/ascii2d_loader.py @ f25328a5

"""
ASCII 2D Loader
"""
from sas.sascalc.dataloader.data_info import Data2D
from sas.sascalc.file_converter.nxcansas_writer import NXcanSASWriter
import numpy as np

# ISIS 2D ASCII File Format
# line: property
# 0: File header
# 1: q_x axis label and units
# 2: q_y axis label and units
# 3: Intensity axis label and units
# 4: nUseRec - number of lines of user content following this line
# 5 - 5+nUseRec: user content
# Number of qx points
# List of qx points
# Number of qy points
# List of qy points
# numqx numqy scale

class ASCII2DLoader(object):

    def __init__(self, data_path):
        self.data_path = data_path

    def load(self):
        file_handle = open(self.data_path, 'r')
        file_buffer = file_handle.read()
        all_lines = file_buffer.splitlines()

        def _load_qs(lines, start_line, num_points):
            qs = np.zeros(num_points)
            n = start_line
            filled = 0
            while filled < num_points:
                row = np.fromstring(lines[n], dtype=np.float32, sep=' ')
                qs[filled:filled+len(row)] = row
                filled += len(row)
                n += 1
            return n, qs

        # Skip nUseRec lines
        current_line = 4
        try:
            nUseRec = int(all_lines[current_line].strip()[0])
            current_line += nUseRec + 1
            # Read qx data
            num_qs = int(all_lines[current_line].strip())
            current_line += 1
            current_line, qx = _load_qs(all_lines, current_line, num_qs)

            # Read qy data
            num_qs = int(all_lines[current_line].strip())
            current_line += 1
            current_line, qy = _load_qs(all_lines, current_line, num_qs)
        except ValueError as e:
            err_msg = "File incorrectly formatted.\n"
            if str(e).find('broadcast') != -1:
                err_msg += "Incorrect number of q data points provided. "
                err_msg += "Expected {}.".format(num_qs)
            elif str(e).find('invalid literal') != -1:
                err_msg += "Expected integer on line {}. Instead got '{}'".format(current_line + 1,
                    all_lines[current_line])
            else:
                err_msg += str(e)
            raise ValueError(err_msg)

        # dimensions: [width, height, scale]
        try:
            dimensions = np.fromstring(all_lines[current_line], dtype=np.float32, sep=' ')
            if len(dimensions) != 3: raise ValueError()
            width = int(dimensions[0])
            height = int(dimensions[1])
        except ValueError as e:
            err_msg = "File incorrectly formatted.\n"
            err_msg += "Expected line {} to be of the form: <num_qx> <num_qy> <scale>.".format(current_line + 1)
            err_msg += " Instead got '{}'.".format(all_lines[current_line])
            raise ValueError(err_msg)

        if width > len(qx) or height > len(qy):
            err_msg = "File incorrectly formatted.\n"
            err_msg += ("Line {} says to use {}x{} points. "
                "Only {}x{} provided.").format(current_line + 1, width, height,
                len(qx), len(qy))
            raise ValueError(err_msg)

        # More qx and/or qy points can be provided than are actually used
        qx = qx[:width]
        qy = qy[:height]

        current_line += 1
        # iflag = 1 => Only intensity data (not dealt with here)
        # iflag = 2 => q axis and intensity data
        # iflag = 3 => q axis, intensity and error data
        try:
            iflag = int(all_lines[current_line].strip()[0])
            if iflag <= 0 or iflag > 3: raise ValueError()
        except:
            err_msg = "File incorrectly formatted.\n"
            iflag = all_lines[current_line].strip()[0]
            err_msg += "Expected iflag on line {} to be 1, 2 or 3. Instead got '{}'.".format(current_line+1, iflag)
            raise ValueError(err_msg)

        current_line += 1

        current_line, I = _load_qs(all_lines, current_line, width*height)
        dI = np.zeros(width*height)

        if iflag == 3:
            _, dI = _load_qs(all_lines, current_line, width*height)

        # Format data for use with Data2D
        qx = list(qx) * height
        qy = np.array([[y] * width for y in qy]).flatten()

        data = Data2D(qx_data=qx, qy_data=qy, data=I, err_data=dI)

        return data