<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Functions for supernova light curve I/O"""

from __future__ import print_function

from warnings import warn
import os
import sys
import re
import json

import numpy as np
from astropy.utils import OrderedDict as odict
from astropy.table import Table
from astropy.io import fits
from astropy import wcs
from astropy.extern import six

from .photdata import dict_to_array

__all__ = ['read_lc', 'write_lc', 'load_example_data', 'read_griddata_ascii',
           'read_griddata_fits', 'write_griddata_ascii', 'write_griddata_fits',
           'read_3dgriddata_fits']


def _stripcomment(line, char='#'):
    pos = line.find(char)
    if pos == -1:
        return line
    else:
        return line[:pos]


def _cast_str(s):
    try:
        return int(s)
    except:
        try:
            return float(s)
        except:
            return s.strip()


def read_griddata_ascii(name_or_obj):
    """Read 2-d grid data from a text file.

    Each line has values `x0 x1 y`. Space separated.
    x1 values are only read for first x0 value. Others are assumed
    to match.

    Parameters
    ----------
    name_or_obj : str or file-like object

    Returns
    -------
    x0 : numpy.ndarray
        1-d array.
    x1 : numpy.ndarray
        1-d array.
    y : numpy.ndarray
        2-d array of shape (len(x0), len(x1)).
    """

    if isinstance(name_or_obj, six.string_types):
        f = open(name_or_obj, 'r')
    else:
        f = name_or_obj

    x0 = []    # x0 values.
    x1 = None  # x1 values for first x0 value, assume others are the same.
    y = []     # 2-d array of internal values

    x0_current = None
    x1_current = []
    y1_current = []
    for line in f:
        stripped_line = _stripcomment(line)
        if len(stripped_line) == 0:
            continue
        x0_tmp, x1_tmp, y_tmp = map(float, stripped_line.split())
        if x0_current is None:
            x0_current = x0_tmp  # Initialize first time

        # If there is a new x0 value, ingest the old one and reset values
        if x0_tmp != x0_current:
            x0.append(x0_current)
            if x1 is None:
                x1 = x1_current
            y.append(y1_current)

            x0_current = x0_tmp
            x1_current = []
            y1_current = []

        x1_current.append(x1_tmp)
        y1_current.append(y_tmp)

    # Ingest the last x0 value and y1 array
    x0.append(x0_current)
    y.append(y1_current)

    f.close()
    return np.array(x0), np.array(x1), np.array(y)


def read_3dgriddata_fits(name_or_obj, ext=0):
    """Read a 3-d grid of data from a FITS file, where the grid coordinates
    are encoded in the FITS-WCS header keywords.

    Parameters
    ----------
    name_or_obj : str or file-like object

    Returns
    -------
    x0 : numpy.ndarray
        1-d array.
    x1 : numpy.ndarray
        1-d array.
    x2 : numpy.ndarray
        1-d array.
    y : numpy.ndarray
        3-d array of shape (len(x0), len(x1), len(x2)).
    """

    hdulist = fits.open(name_or_obj)
    w = wcs.WCS(hdulist[ext].header)
    y = hdulist[ext].data
    nx0, nx1, nx2 = y.shape

    # get x0 values
    coords = np.empty((nx2, 3), dtype=np.float32)
    coords[:, 1:3] = 0.
    coords[:, 0] = np.arange(nx2)  # x2 = FITS AXIS1 ("z" coordinates)
    x2 = w.wcs_pix2world(coords, 0)[:, 0]

    # get x1 values
    coords = np.empty((nx1, 3), dtype=np.float32)
    coords[:, 0] = 0.
    coords[:, 1] = np.arange(nx1)  # x1 = FITS AXIS2 ("y" coordinates)
    coords[:, 2] = 0.
    x1 = w.wcs_pix2world(coords, 0)[:, 1]

    # get x2 values
    coords = np.empty((nx0, 3), dtype=np.float32)
    coords[:, 2] = np.arange(nx0)  # x0 = FITS AXIS3 ("x" coordinates)
    coords[:, 0:2] = 0.
    x0 = w.wcs_pix2world(coords, 0)[:, 2]

    hdulist.close()

    return x0, x1, x2, y


def read_griddata_fits(name_or_obj, ext=0):
    """Read a 2-d grid of data from a FITS file, where the grid coordinates
    are encoded in the FITS-WCS header keywords.

    Parameters
    ----------
    name_or_obj : str or file-like object

    Returns
    -------
    x0 : numpy.ndarray
        1-d array.
    x1 : numpy.ndarray
        1-d array.
    y : numpy.ndarray
        2-d array of shape (len(x0), len(x1)).
    """

    hdulist = fits.open(name_or_obj)
    w = wcs.WCS(hdulist[ext].header)
    y = hdulist[ext].data
    nx0, nx1 = y.shape

    # get x0 values
    coords = np.empty((nx0, 2), dtype=np.float32)
    coords[:, 0] = 0.
    coords[:, 1] = np.arange(nx0)  # x0 = FITS AXIS2 ("y" coordinates)
    x0 = w.wcs_pix2world(coords, 0)[:, 1]

    # get x1 values
    coords = np.empty((nx1, 2), dtype=np.float32)
    coords[:, 0] = np.arange(nx1)  # x1 = FITS AXIS1 ("x" coordinates)
    coords[:, 1] = 0.
    x1 = w.wcs_pix2world(coords, 0)[:, 0]

    hdulist.close()

    return x0, x1, y


def write_griddata_ascii(x0, x1, y, name_or_obj):
    """Write 2-d grid data to a text file.

    Each line has values `x0 x1 y`. Space separated.

    Parameters
    ----------
    x0 : numpy.ndarray
        1-d array.
    x1 : numpy.ndarray
        1-d array.
    y : numpy.ndarray
        2-d array of shape (len(x0), len(x1)).
    name_or_obj : str or file-like object
        Filename to write to or open file.
    """

    if isinstance(name_or_obj, six.string_types):
        f = open(name_or_obj, 'r')
    else:
        f = name_or_obj

    for j in range(len(x0)):
        for i in range(len(x1)):
            f.write("{0:.7g} {1:.7g} {2:.7g}\n".format(x0[j], x1[i], y[j, i]))
<commit_msg>make read_griddata_fits() multi-dimensional<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Functions for supernova light curve I/O"""

from __future__ import print_function

from warnings import warn
import os
import sys
import re
import json

import numpy as np
from astropy.utils import OrderedDict as odict
from astropy.table import Table
from astropy.io import fits
from astropy import wcs
from astropy.extern import six

from .photdata import dict_to_array

__all__ = ['read_lc', 'write_lc', 'load_example_data', 'read_griddata_ascii',
           'read_griddata_fits', 'write_griddata_ascii', 'write_griddata_fits',
           'read_3dgriddata_fits']


def _stripcomment(line, char='#'):
    pos = line.find(char)
    if pos == -1:
        return line
    else:
        return line[:pos]


def _cast_str(s):
    try:
        return int(s)
    except:
        try:
            return float(s)
        except:
            return s.strip()


def read_griddata_ascii(name_or_obj):
    """Read 2-d grid data from a text file.

    Each line has values `x0 x1 y`. Space separated.
    x1 values are only read for first x0 value. Others are assumed
    to match.

    Parameters
    ----------
    name_or_obj : str or file-like object

    Returns
    -------
    x0 : numpy.ndarray
        1-d array.
    x1 : numpy.ndarray
        1-d array.
    y : numpy.ndarray
        2-d array of shape (len(x0), len(x1)).
    """

    if isinstance(name_or_obj, six.string_types):
        f = open(name_or_obj, 'r')
    else:
        f = name_or_obj

    x0 = []    # x0 values.
    x1 = None  # x1 values for first x0 value, assume others are the same.
    y = []     # 2-d array of internal values

    x0_current = None
    x1_current = []
    y1_current = []
    for line in f:
        stripped_line = _stripcomment(line)
        if len(stripped_line) == 0:
            continue
        x0_tmp, x1_tmp, y_tmp = map(float, stripped_line.split())
        if x0_current is None:
            x0_current = x0_tmp  # Initialize first time

        # If there is a new x0 value, ingest the old one and reset values
        if x0_tmp != x0_current:
            x0.append(x0_current)
            if x1 is None:
                x1 = x1_current
            y.append(y1_current)

            x0_current = x0_tmp
            x1_current = []
            y1_current = []

        x1_current.append(x1_tmp)
        y1_current.append(y_tmp)

    # Ingest the last x0 value and y1 array
    x0.append(x0_current)
    y.append(y1_current)

    f.close()
    return np.array(x0), np.array(x1), np.array(y)


def read_3dgriddata_fits(name_or_obj, ext=0):
    """Read a 3-d grid of data from a FITS file, where the grid coordinates
    are encoded in the FITS-WCS header keywords.

    Parameters
    ----------
    name_or_obj : str or file-like object

    Returns
    -------
    x0 : numpy.ndarray
        1-d array.
    x1 : numpy.ndarray
        1-d array.
    x2 : numpy.ndarray
        1-d array.
    y : numpy.ndarray
        3-d array of shape (len(x0), len(x1), len(x2)).
    """

    hdulist = fits.open(name_or_obj)
    w = wcs.WCS(hdulist[ext].header)
    y = hdulist[ext].data
    nx0, nx1, nx2 = y.shape

    # get x0 values
    coords = np.empty((nx2, 3), dtype=np.float32)
    coords[:, 1:3] = 0.
    coords[:, 0] = np.arange(nx2)  # x2 = FITS AXIS1 ("z" coordinates)
    x2 = w.wcs_pix2world(coords, 0)[:, 0]

    # get x1 values
    coords = np.empty((nx1, 3), dtype=np.float32)
    coords[:, 0] = 0.
    coords[:, 1] = np.arange(nx1)  # x1 = FITS AXIS2 ("y" coordinates)
    coords[:, 2] = 0.
    x1 = w.wcs_pix2world(coords, 0)[:, 1]

    # get x2 values
    coords = np.empty((nx0, 3), dtype=np.float32)
    coords[:, 2] = np.arange(nx0)  # x0 = FITS AXIS3 ("x" coordinates)
    coords[:, 0:2] = 0.
    x0 = w.wcs_pix2world(coords, 0)[:, 2]

    hdulist.close()

    return x0, x1, x2, y


def read_griddata_fits(name_or_obj, ext=0):
    """Read a 2-d grid of data from a FITS file, where the grid coordinates
    are encoded in the FITS-WCS header keywords.

    Parameters
    ----------
    name_or_obj : str or file-like object

    Returns
    -------
    x0, x1, ... : `~numpy.ndarray`s
        1-d arrays giving coordinates of grid. The number of these arrays will
        depend on the dimension of the data array. For example, if the data have
        two dimensions, a total of three arrays will be returned: ``x0, x1, y``,
        with ``x0`` giving the coordinates of the first axis of ``y``. If the
        data have three dimensions, a total of four arrays will be returned:
        ``x0, x1, x2, y``, and so on with higher dimensions.
    y : `~numpy.ndarray`
        n-d array of shape ``(len(x0), len(x1), ...)``. For three dimensions
        for example, the value at ``y[i, j, k]`` corresponds to coordinates
        ``(x0[i], x1[j], x2[k])``.
    """

    hdulist = fits.open(name_or_obj)
    w = wcs.WCS(hdulist[ext].header)
    y = hdulist[ext].data

    # get abcissa values (coordinates at grid values)
    xs = []
    for i in range(y.ndim):
        j = y.ndim - i  # The i-th axis (in Python) corresponds to FITS AXISj
        coords = np.zeros((y.shape[i], y.ndim), dtype=np.float32)
        coords[:, j-1] = np.arange(nx)
        x = w.wcs_pix2world(coords, 0)[:, j-1]
        xs.append(x)

    hdulist.close()

    return tuple(xs) + (y,)


def write_griddata_ascii(x0, x1, y, name_or_obj):
    """Write 2-d grid data to a text file.

    Each line has values `x0 x1 y`. Space separated.

    Parameters
    ----------
    x0 : numpy.ndarray
        1-d array.
    x1 : numpy.ndarray
        1-d array.
    y : numpy.ndarray
        2-d array of shape (len(x0), len(x1)).
    name_or_obj : str or file-like object
        Filename to write to or open file.
    """

    if isinstance(name_or_obj, six.string_types):
        f = open(name_or_obj, 'r')
    else:
        f = name_or_obj

    for j in range(len(x0)):
        for i in range(len(x1)):
            f.write("{0:.7g} {1:.7g} {2:.7g}\n".format(x0[j], x1[i], y[j, i]))

{'fraction_non_alphanumeric': 0.12539641724522158, 'fraction_numerical': 0.030084854718436617, 'mean_word_length': 2.1594909287841864, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '709913', 'n_tokens_mistral': 4397, 'n_tokens_neox': 4238, 'n_words': 1222}