text:
# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
#
#   See COPYING file distributed along with the NiBabel package for the
#   copyright and license terms.
#
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
''' Read / write access to NIfTI1 image format

NIfTI1 format defined at http://nifti.nimh.nih.gov/nifti-1/
'''
from __future__ import division, print_function
import warnings
from io import BytesIO
from six import string_types

import numpy as np
import numpy.linalg as npl

from .py3k import asstr
from .volumeutils import Recoder, make_dt_codes, endian_codes
from .spatialimages import HeaderDataError, ImageFileError
from .batteryrunners import Report
from .quaternions import fillpositive, quat2mat, mat2quat
from . import analyze  # module import
from .spm99analyze import SpmAnalyzeHeader
from .casting import have_binary128
from .pydicom_compat import have_dicom, pydicom as pdcm
from .testing import setup_test  # noqa

# nifti1 flat header definition for Analyze-like first 348 bytes
# first number in comments indicates offset in file header in bytes
header_dtd = [
    ('sizeof_hdr', 'i4'),      # 0; must be 348
    ('data_type', 'S10'),      # 4; unused
    ('db_name', 'S18'),        # 14; unused
    ('extents', 'i4'),         # 32; unused
    ('session_error', 'i2'),   # 36; unused
    ('regular', 'S1'),         # 38; unused
    ('dim_info', 'u1'),        # 39; MRI slice ordering code
    ('dim', 'i2', (8,)),       # 40; data array dimensions
    ('intent_p1', 'f4'),       # 56; first intent parameter
    ('intent_p2', 'f4'),       # 60; second intent parameter
    ('intent_p3', 'f4'),       # 64; third intent parameter
    ('intent_code', 'i2'),     # 68; NIFTI intent code
    ('datatype', 'i2'),        # 70; it's the datatype
    ('bitpix', 'i2'),          # 72; number of bits per voxel
    ('slice_start', 'i2'),     # 74; first slice index
    ('pixdim', 'f4', (8,)),    # 76; grid spacings (units below)
    ('vox_offset', 'f4'),      # 108; offset to data in image file
    ('scl_slope', 'f4'),       # 112; data scaling slope
    ('scl_inter', 'f4'),       # 116; data scaling intercept
    ('slice_end', 'i2'),       # 120; last slice index
    ('slice_code', 'u1'),      # 122; slice timing order
    ('xyzt_units', 'u1'),      # 123; units of pixdim[1..4]
    ('cal_max', 'f4'),         # 124; max display intensity
    ('cal_min', 'f4'),         # 128; min display intensity
    ('slice_duration', 'f4'),  # 132; time for 1 slice
    ('toffset', 'f4'),         # 136; time axis shift
    ('glmax', 'i4'),           # 140; unused
    ('glmin', 'i4'),           # 144; unused
    ('descrip', 'S80'),        # 148; any text
    ('aux_file', 'S24'),       # 228; auxiliary filename
    ('qform_code', 'i2'),      # 252; xform code
    ('sform_code', 'i2'),      # 254; xform code
    ('quatern_b', 'f4'),       # 256; quaternion b param
    ('quatern_c', 'f4'),       # 260; quaternion c param
    ('quatern_d', 'f4'),       # 264; quaternion d param
    ('qoffset_x', 'f4'),       # 268; quaternion x shift
    ('qoffset_y', 'f4'),       # 272; quaternion y shift
    ('qoffset_z', 'f4'),       # 276; quaternion z shift
    ('srow_x', 'f4', (4,)),    # 280; 1st row affine transform
    ('srow_y', 'f4', (4,)),    # 296; 2nd row affine transform
    ('srow_z', 'f4', (4,)),    # 312; 3rd row affine transform
    ('intent_name', 'S16'),    # 328; name or meaning of data
    ('magic', 'S4')            # 344; must be 'ni1\0' or 'n+1\0'
]

# Full header numpy dtype
header_dtype = np.dtype(header_dtd)

# datatypes not in analyze format, with codes
if have_binary128():
    # Only enable 128 bit floats if we really have IEEE binary 128 longdoubles
    _float128t = np.longdouble
    _complex256t = np.longcomplex
else:
    _float128t = np.void
    _complex256t = np.void

_dtdefs = (  # code, label, dtype definition, niistring
    (0, 'none', np.void, ""),
    (1, 'binary', np.void, ""),
    (2, 'uint8', np.uint8, "NIFTI_TYPE_UINT8"),
    (4, 'int16', np.int16, "NIFTI_TYPE_INT16"),
    (8, 'int32', np.int32, "NIFTI_TYPE_INT32"),
    (16, 'float32', np.float32, "NIFTI_TYPE_FLOAT32"),
    (32, 'complex64', np.complex64, "NIFTI_TYPE_COMPLEX64"),
    (64, 'float64', np.float64, "NIFTI_TYPE_FLOAT64"),
    (128, 'RGB', np.dtype([('R', 'u1'),
                           ('G', 'u1'),
                           ('B', 'u1')]), "NIFTI_TYPE_RGB24"),
    (255, 'all', np.void, ''),
    (256, 'int8', np.int8, "NIFTI_TYPE_INT8"),
    (512, 'uint16', np.uint16, "NIFTI_TYPE_UINT16"),
    (768, 'uint32', np.uint32, "NIFTI_TYPE_UINT32"),
    (1024, 'int64', np.int64, "NIFTI_TYPE_INT64"),
    (1280, 'uint64', np.uint64, "NIFTI_TYPE_UINT64"),
    (1536, 'float128', _float128t, "NIFTI_TYPE_FLOAT128"),
    (1792, 'complex128', np.complex128, "NIFTI_TYPE_COMPLEX128"),
    (2048, 'complex256', _complex256t, "NIFTI_TYPE_COMPLEX256"),
    (2304, 'RGBA', np.dtype([('R', 'u1'),
                             ('G', 'u1'),
                             ('B', 'u1'),
                             ('A', 'u1')]), "NIFTI_TYPE_RGBA32"),
)

# Make full code alias bank, including dtype column
data_type_codes = make_dt_codes(_dtdefs)

# Transform (qform, sform) codes
xform_codes = Recoder((  # code, label, niistring
    (0, 'unknown', "NIFTI_XFORM_UNKNOWN"),
    (1, 'scanner', "NIFTI_XFORM_SCANNER_ANAT"),
    (2, 'aligned', "NIFTI_XFORM_ALIGNED_ANAT"),
    (3, 'talairach', "NIFTI_XFORM_TALAIRACH"),
    (4, 'mni', "NIFTI_XFORM_MNI_152")), fields=('code', 'label', 'niistring'))

# unit codes
unit_codes = Recoder((  # code, label
    (0, 'unknown'),
    (1, 'meter'),
    (2, 'mm'),
    (3, 'micron'),
    (8, 'sec'),
    (16, 'msec'),
    (24, 'usec'),
    (32, 'hz'),
    (40, 'ppm'),
    (48, 'rads')), fields=('code', 'label'))

slice_order_codes = Recoder((  # code, label
    (0, 'unknown'),
    (1, 'sequential increasing', 'seq inc'),
    (2, 'sequential decreasing', 'seq dec'),
    (3, 'alternating increasing', 'alt inc'),
    (4, 'alternating decreasing', 'alt dec'),
    (5, 'alternating increasing 2', 'alt inc 2'),
    (6, 'alternating decreasing 2', 'alt dec 2')), fields=('code', 'label'))

intent_codes = Recoder((
    # code, label, parameters description tuple
    (0, 'none', (), "NIFTI_INTENT_NONE"),
    (2, 'correlation', ('p1 = DOF',), "NIFTI_INTENT_CORREL"),
    (3, 't test', ('p1 = DOF',), "NIFTI_INTENT_TTEST"),
    (4, 'f test', ('p1 = numerator DOF', 'p2 = denominator DOF'),
     "NIFTI_INTENT_FTEST"),
    (5, 'z score', (), "NIFTI_INTENT_ZSCORE"),
    (6, 'chi2', ('p1 = DOF',), "NIFTI_INTENT_CHISQ"),
    # two parameter beta distribution
    (7, 'beta',
     ('p1=a', 'p2=b'),
     "NIFTI_INTENT_BETA"),
    # Prob(x) = (p1 choose x) * p2^x * (1-p2)^(p1-x), for x=0,1,...,p1
    (8, 'binomial',
     ('p1 = number of trials', 'p2 = probability per trial'),
     "NIFTI_INTENT_BINOM"),
    # 2 parameter gamma
    # Density(x) proportional to  # x^(p1-1) * exp(-p2*x)
    (9, 'gamma',
     ('p1 = shape, p2 = scale', 2),
     "NIFTI_INTENT_GAMMA"),
    (10, 'poisson',
     ('p1 = mean',),
     "NIFTI_INTENT_POISSON"),
    (11, 'normal',
     ('p1 = mean', 'p2 = standard deviation',),
     "NIFTI_INTENT_NORMAL"),
    (12, 'non central f test',
     ('p1 = numerator DOF',
      'p2 = denominator DOF',
      'p3 = numerator noncentrality parameter',),
     "NIFTI_INTENT_FTEST_NONC"),
    (13, 'non central chi2',
     ('p1 = DOF', 'p2 = noncentrality parameter',),
     "NIFTI_INTENT_CHISQ_NONC"),
    (14, 'logistic',
     ('p1 = location', 'p2 = scale',),
     "NIFTI_INTENT_LOGISTIC"),
    (15, 'laplace',
     ('p1 = location', 'p2 = scale'),
     "NIFTI_INTENT_LAPLACE"),
    (16, 'uniform',
     ('p1 = lower end', 'p2 = upper end'),
     "NIFTI_INTENT_UNIFORM"),
    (17, 'non central t test',
     ('p1 = DOF', 'p2 = noncentrality parameter'),
     "NIFTI_INTENT_TTEST_NONC"),
    (18, 'weibull',
     ('p1 = location', 'p2 = scale, p3 = power'),
     "NIFTI_INTENT_WEIBULL"),
    # p1 = 1 = 'half normal' distribution
    # p1 = 2 = Rayleigh distribution
    # p1 = 3 = Maxwell-Boltzmann distribution.
    (19, 'chi', ('p1 = DOF',), "NIFTI_INTENT_CHI"),
    (20, 'inverse gaussian',
     ('pi = mu', 'p2 = lambda'),
     "NIFTI_INTENT_INVGAUSS"),
    (21, 'extreme value 1',
     ('p1 = location', 'p2 = scale'),
     "NIFTI_INTENT_EXTVAL"),
    (22, 'p value', (), "NIFTI_INTENT_PVAL"),
    (23, 'log p value', (), "NIFTI_INTENT_LOGPVAL"),
    (24, 'log10 p value', (), "NIFTI_INTENT_LOG10PVAL"),
    (1001, 'estimate', (), "NIFTI_INTENT_ESTIMATE"),
    (1002, 'label', (), "NIFTI_INTENT_LABEL"),
    (1003, 'neuroname', (), "NIFTI_INTENT_NEURONAME"),
    (1004, 'general matrix',
     ('p1 = M', 'p2 = N'),
     "NIFTI_INTENT_GENMATRIX"),
    (1005, 'symmetric matrix', ('p1 = M',), "NIFTI_INTENT_SYMMATRIX"),
    (1006, 'displacement vector', (), "NIFTI_INTENT_DISPVECT"),
    (1007, 'vector', (), "NIFTI_INTENT_VECTOR"),
    (1008, 'pointset', (), "NIFTI_INTENT_POINTSET"),
    (1009, 'triangle', (), "NIFTI_INTENT_TRIANGLE"),
    (1010, 'quaternion', (), "NIFTI_INTENT_QUATERNION"),
    (1011, 'dimensionless', (), "NIFTI_INTENT_DIMLESS"),
    (2001, 'time series',
     (),
     "NIFTI_INTENT_TIME_SERIES",
     "NIFTI_INTENT_TIMESERIES"),  # this mis-spell occurs in the wild
    (2002, 'node index', (), "NIFTI_INTENT_NODE_INDEX"),
    (2003, 'rgb vector', (), "NIFTI_INTENT_RGB_VECTOR"),
    (2004, 'rgba vector', (), "NIFTI_INTENT_RGBA_VECTOR"),
    (2005, 'shape', (), "NIFTI_INTENT_SHAPE"),
    # FSL-specific intent codes - codes used by FNIRT
    # ($FSLDIR/warpfns/fnirt_file_reader.h:104)
    (2006, 'fnirt disp field', (), 'FSL_FNIRT_DISPLACEMENT_FIELD'),
    (2007, 'fnirt cubic spline coef', (), 'FSL_CUBIC_SPLINE_COEFFICIENTS'),
    (2008, 'fnirt dct coef', (), 'FSL_DCT_COEFFICIENTS'),
    (2009, 'fnirt quad spline coef', (), 'FSL_QUADRATIC_SPLINE_COEFFICIENTS'),
    # FSL-specific intent codes - codes used by TOPUP
    # ($FSLDIR/topup/topup_file_io.h:104)
    (2016, 'topup cubic spline coef ', (),
     'FSL_TOPUP_CUBIC_SPLINE_COEFFICIENTS'),
    (2017, 'topup quad spline coef', (),
     'FSL_TOPUP_QUADRATIC_SPLINE_COEFFICIENTS'),
    (2018, 'topup field', (), 'FSL_TOPUP_FIELD'),
), fields=('code', 'label', 'parameters', 'niistring'))

class Nifti1Extension(object):
    """Baseclass for NIfTI1 header extensions.

    This class is sufficient to handle very simple text-based extensions, such
    as `comment`. More sophisticated extensions should/will be supported by
    dedicated subclasses.
    """

    def __init__(self, code, content):
        """
        Parameters
        ----------
        code : int or str
          Canonical extension code as defined in the NIfTI standard, given
          either as integer or corresponding label
          (see :data:`~nibabel.nifti1.extension_codes`)
        content : str
          Extension content as read from the NIfTI file header. This content is
          converted into a runtime representation.
        """
        try:
            self._code = extension_codes.code[code]
        except KeyError:
            # XXX or fail or at least complain?
            self._code = code
        self._content = self._unmangle(content)

    def _unmangle(self, value):
        """Convert the extension content into its runtime representation.

        The default implementation does nothing at all.

        Parameters
        ----------
        value : str
          Extension content as read from file.

        Returns
        -------
        The same object that was passed as `value`.

        Notes
        -----
        Subclasses should reimplement this method to provide the desired
        unmangling procedure and may return any type of object.
        """
        return value

    def _mangle(self, value):
        """Convert the extension content into NIfTI file header representation.

        The default implementation does nothing at all.

        Parameters
        ----------
        value : str
          Extension content in runtime form.

        Returns
        -------
        str

        Notes
        -----
        Subclasses should reimplement this method to provide the desired
        mangling procedure.
        """
        return value

    def get_code(self):
        """Return the canonical extension type code."""
        return self._code

    def get_content(self):
        """Return the extension content in its runtime representation."""
        return self._content

    def get_sizeondisk(self):
        """Return the size of the extension in the NIfTI file.
        """
        # need raw value size plus 8 bytes for esize and ecode
        size = len(self._mangle(self._content))
        size += 8
        # extensions size has to be a multiple of 16 bytes
        if size % 16 != 0:
            size += 16 - (size % 16)
        return size

    def __repr__(self):
        try:
            code = extension_codes.label[self._code]
        except KeyError:
            # deal with unknown codes
            code = self._code

        s = "Nifti1Extension('%s', '%s')" % (code, self._content)
        return s

    def __eq__(self, other):
        return (self._code, self._content) == (other._code, other._content)

    def __ne__(self, other):
        return not self == other

    def write_to(self, fileobj, byteswap):
        ''' Write header extensions to fileobj

        Write starts at fileobj current file position.

        Parameters
        ----------
        fileobj : file-like object
           Should implement ``write`` method
        byteswap : boolean
          Flag if byteswapping the data is required.

        Returns
        -------
        None
        '''
        extstart = fileobj.tell()
        rawsize = self.get_sizeondisk()
        # write esize and ecode first
        extinfo = np.array((rawsize, self._code), dtype=np.int32)
        if byteswap:
            extinfo = extinfo.byteswap()
        fileobj.write(extinfo.tostring())
        # followed by the actual extension content
        # XXX if mangling upon load is implemented, it should be reverted here
        fileobj.write(self._mangle(self._content))
        # be nice and zero out remaining part of the extension till the
        # next 16 byte border
        fileobj.write(b'\x00' * (extstart + rawsize - fileobj.tell()))

class Nifti1DicomExtension(Nifti1Extension):
    """NIfTI1 DICOM header extension

    This class is a thin wrapper around pydicom to read a binary DICOM
    byte string. If pydicom is available, content is exposed as a Dicom Dataset.
    Otherwise, this silently falls back to the standard NiftiExtension class
    and content is the raw bytestring loaded directly from the nifti file
    header.
    """
    def __init__(self, code, content, parent_hdr=None):
        """
        Parameters
        ----------
        code : int or str
          Canonical extension code as defined in the NIfTI standard, given
          either as integer or corresponding label
          (see :data:`~nibabel.nifti1.extension_codes`)
        content : bytes or pydicom Dataset or None
          Extension content - either a bytestring as read from the NIfTI file
          header or an existing pydicom Dataset. If a bystestring, the content
          is converted into a Dataset on initialization. If None, a new empty
          Dataset is created.
        parent_hdr : :class:`~nibabel.nifti1.Nifti1Header`, optional
          If a dicom extension belongs to an existing
          :class:`~nibabel.nifti1.Nifti1Header`, it may be provided here to
          ensure that the DICOM dataset is written with correctly corresponding
          endianness; otherwise it is assumed the dataset is little endian.

        Notes
        -----

        code should always be 2 for DICOM.
        """

        self._code = code
        if parent_hdr:
            self._is_little_endian = parent_hdr.endianness == '<'
        else:
            self._is_little_endian = True
        if isinstance(content, pdcm.dataset.Dataset):
            self._is_implicit_VR = False
            self._raw_content = self._mangle(content)
            self._content = content
        elif isinstance(content, bytes):  # Got a byte string - unmangle it
            self._raw_content = content
            self._is_implicit_VR = self._guess_implicit_VR()
            ds = self._unmangle(content, self._is_implicit_VR,
                                self._is_little_endian)
            self._content = ds
        elif content is None:  # initialize a new dicom dataset
            self._is_implicit_VR = False
            self._content = pdcm.dataset.Dataset()
        else:
            raise TypeError("content must be either a bytestring or a pydicom "
                            "Dataset. Got %s" % content.__class__)

    def _guess_implicit_VR(self):
        """Try to guess DICOM syntax by checking for valid VRs.

        Without a DICOM Transfer Syntax, it's difficult to tell if Value
        Representations (VRs) are included in the DICOM encoding or not.
        This reads where the first VR would be and checks it against a list of
        valid VRs
        """
        potential_vr = self._raw_content[4:6].decode()
        if potential_vr in pdcm.values.converters.keys():
            implicit_VR = False
        else:
            implicit_VR = True
        return implicit_VR

    def _unmangle(self, value, is_implicit_VR=False, is_little_endian=True):
        bio = BytesIO(value)
        ds = pdcm.filereader.read_dataset(bio,
                                          is_implicit_VR,
                                          is_little_endian)
        return ds

    def _mangle(self, dataset):
        bio = BytesIO()
        dio = pdcm.filebase.DicomFileLike(bio)
        dio.is_implicit_VR = self._is_implicit_VR
        dio.is_little_endian = self._is_little_endian
        ds_len = pdcm.filewriter.write_dataset(dio, dataset)
        dio.seek(0)
        return dio.read(ds_len)

# NIfTI header extension type codes (ECODE)
# see nifti1_io.h for a complete list of all known extensions and
# references to their description or contacts of the respective
# initiators
extension_codes = Recoder((
    (0, "ignore", Nifti1Extension),
    (2, "dicom", Nifti1DicomExtension if have_dicom else Nifti1Extension),
    (4, "afni", Nifti1Extension),
    (6, "comment", Nifti1Extension),
    (8, "xcede", Nifti1Extension),
    (10, "jimdiminfo", Nifti1Extension),
    (12, "workflow_fwds", Nifti1Extension),
    (14, "freesurfer", Nifti1Extension),
    (16, "pypickle", Nifti1Extension),
), fields=('code', 'label', 'handler'))

class Nifti1Extensions(list):
    """Simple extension collection, implemented as a list-subclass.
    """

    def count(self, ecode):
        """Returns the number of extensions matching a given *ecode*.

        Parameters
        ----------
        code : int | str
            The ecode can be specified either literal or as numerical value.
        """
        count = 0
        code = extension_codes.code[ecode]
        for e in self:
            if e.get_code() == code:
                count += 1
        return count

    def get_codes(self):
        """Return a list of the extension code of all available extensions"""
        return [e.get_code() for e in self]

    def get_sizeondisk(self):
        """Return the size of the complete header extensions in the NIfTI file.
        """
        return np.sum([e.get_sizeondisk() for e in self])

    def __repr__(self):
        s = "Nifti1Extensions(%s)" % ', '.join(str(e) for e in self)
        return s

    def __cmp__(self, other):
        return cmp(list(self), list(other))

    def write_to(self, fileobj, byteswap):
        ''' Write header extensions to fileobj

        Write starts at fileobj current file position.

        Parameters
        ----------
        fileobj : file-like object
           Should implement ``write`` method
        byteswap : boolean
          Flag if byteswapping the data is required.

        Returns
        -------
        None
        '''
        for e in self:
            e.write_to(fileobj, byteswap)

    @classmethod
    def from_fileobj(klass, fileobj, size, byteswap):
        '''Read header extensions from a fileobj

        Parameters
        ----------
        fileobj : file-like object
            We begin reading the extensions at the current file position
        size : int
            Number of bytes to read. If negative, fileobj will be read till its
            end.
        byteswap : boolean
            Flag if byteswapping the read data is required.

        Returns
        -------
        An extension list. This list might be empty in case not extensions
        were present in fileobj.
        '''
        # make empty extension list
        extensions = klass()
        # assume the file pointer is at the beginning of any extensions.
        # read until the whole header is parsed (each extension is a multiple
        # of 16 bytes) or in case of a separate header file till the end
        # (break inside the body)
        while size >= 16 or size < 0:
            # the next 8 bytes should have esize and ecode
            ext_def = fileobj.read(8)
            # nothing was read and instructed to read till the end
            # -> assume all extensions where parsed and break
            if not len(ext_def) and size < 0:
                break
            # otherwise there should be a full extension header
            if not len(ext_def) == 8:
                raise HeaderDataError('failed to read extension header')
            ext_def = np.frombuffer(ext_def, dtype=np.int32)
            if byteswap:
                ext_def = ext_def.byteswap()
            # be extra verbose
            ecode = ext_def[1]
            esize = ext_def[0]
            if esize % 16:
                warnings.warn(
                    'Extension size is not a multiple of 16 bytes; '
                    'Assuming size is correct and hoping for the best',
                    UserWarning)
            # read extension itself; esize includes the 8 bytes already read
            evalue = fileobj.read(int(esize - 8))
            if not len(evalue) == esize - 8:
                raise HeaderDataError('failed to read extension content')
            # note that we read a full extension
            size -= esize
            # store raw extension content, but strip trailing NULL chars
            evalue = evalue.rstrip(b'\x00')
            # 'extension_codes' also knows the best implementation to handle
            # a particular extension type
            try:
                ext = extension_codes.handler[ecode](ecode, evalue)
            except KeyError:
                # unknown extension type
                # XXX complain or fail or go with a generic extension
                ext = Nifti1Extension(ecode, evalue)
            extensions.append(ext)
        return extensions

class Nifti1Header(SpmAnalyzeHeader):
    ''' Class for NIfTI1 header

    The NIfTI1 header has many more coded fields than the simpler Analyze
    variants.  NIfTI1 headers also have extensions.

    Nifti allows the header to be a separate file, as part of a nifti image /
    header pair, or to precede the data in a single file.  The object needs to
    know which type it is, in order to manage the voxel offset pointing to the
    data, extension reading, and writing the correct magic string.

    This class handles the header-preceding-data case.
    '''
    # Copies of module level definitions
    template_dtype = header_dtype
    _data_type_codes = data_type_codes

    # fields with recoders for their values
    _field_recoders = {'datatype': data_type_codes,
                       'qform_code': xform_codes,
                       'sform_code': xform_codes,
                       'intent_code': intent_codes,
                       'slice_code': slice_order_codes}

    # data scaling capabilities
    has_data_slope = True
    has_data_intercept = True

    # Extension class; should implement __call__ for construction, and
    # ``from_fileobj`` for reading from file
    exts_klass = Nifti1Extensions

    # Signal whether this is single (header + data) file
    is_single = True

    # Default voxel data offsets for single and pair
    pair_vox_offset = 0
    single_vox_offset = 352

    # Magics for single and pair
    pair_magic = b'ni1'
    single_magic = b'n+1'

    # Quaternion threshold near 0, based on float32 precision
    quaternion_threshold = -np.finfo(np.float32).eps * 3

    def __init__(self,
                 binaryblock=None,
                 endianness=None,
                 check=True,
                 extensions=()):
        ''' Initialize header from binary data block and extensions
        '''
        super(Nifti1Header, self).__init__(binaryblock,
                                           endianness,
                                           check)
        self.extensions = self.exts_klass(extensions)

    def copy(self):
        ''' Return copy of header

        Take reference to extensions as well as copy of header contents
        '''
        return self.__class__(
            self.binaryblock,
            self.endianness,
            False,
            self.extensions)

    @classmethod
    def from_fileobj(klass, fileobj, endianness=None, check=True):
        raw_str = fileobj.read(klass.template_dtype.itemsize)
        hdr = klass(raw_str, endianness, check)
        # Read next 4 bytes to see if we have extensions.  The nifti standard
        # has this as a 4 byte string; if the first value is not zero, then we
        # have extensions.
        extension_status = fileobj.read(4)
        # Need to test *slice* of extension_status to preserve byte string type
        # on Python 3
        if len(extension_status) < 4 or extension_status[0:1] == b'\x00':
            return hdr
        # If this is a detached header file read to end
        if not klass.is_single:
            extsize = -1
        else:  # otherwise read until the beginning of the data
            extsize = hdr._structarr['vox_offset'] - fileobj.tell()
        byteswap = endian_codes['native'] != hdr.endianness
        hdr.extensions = klass.exts_klass.from_fileobj(fileobj, extsize,
                                                       byteswap)
        return hdr

    def write_to(self, fileobj):
        # First check that vox offset is large enough; set if necessary
        if self.is_single:
            vox_offset = self._structarr['vox_offset']
            min_vox_offset = (self.single_vox_offset +
                              self.extensions.get_sizeondisk())
            if vox_offset == 0:  # vox offset unset; set as necessary
                self._structarr['vox_offset'] = min_vox_offset
            elif vox_offset < min_vox_offset:
                raise HeaderDataError(
                    'vox offset set to {0}, but need at least {1}'.format(
                        vox_offset, min_vox_offset))
        super(Nifti1Header, self).write_to(fileobj)
        # Write extensions
        if len(self.extensions) == 0:
            # If single file, write required 0 stream to signal no extensions
            if self.is_single:
                fileobj.write(b'\x00' * 4)
            return
        # Signal there are extensions that follow
        fileobj.write(b'\x01\x00\x00\x00')
        byteswap = endian_codes['native'] != self.endianness
        self.extensions.write_to(fileobj, byteswap)

    def get_best_affine(self):
        ''' Select best of available transforms '''
        hdr = self._structarr
        if hdr['sform_code'] != 0:
            return self.get_sform()
        if hdr['qform_code'] != 0:
            return self.get_qform()
        return self.get_base_affine()

    @classmethod
    def default_structarr(klass, endianness=None):
        ''' Create empty header binary block with given endianness '''
        hdr_data = super(Nifti1Header, klass).default_structarr(endianness)
        if klass.is_single:
            hdr_data['magic'] = klass.single_magic
        else:
            hdr_data['magic'] = klass.pair_magic
        return hdr_data

    @classmethod
    def from_header(klass, header=None, check=True):
        ''' Class method to create header from another header

        Extend Analyze header copy by copying extensions from other Nifti
        types.

        Parameters
        ----------
        header : ``Header`` instance or mapping
           a header of this class, or another class of header for
           conversion to this type
        check : {True, False}
           whether to check header for integrity

        Returns
        -------
        hdr : header instance
           fresh header instance of our own class
        '''
        new_hdr = super(Nifti1Header, klass).from_header(header, check)
        if isinstance(header, Nifti1Header):
            new_hdr.extensions[:] = header.extensions[:]
        return new_hdr

    def get_data_shape(self):
        ''' Get shape of data

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> hdr.get_data_shape()
        (0,)
        >>> hdr.set_data_shape((1,2,3))
        >>> hdr.get_data_shape()
        (1, 2, 3)

        Expanding number of dimensions gets default zooms

        >>> hdr.get_zooms()
        (1.0, 1.0, 1.0)

        Notes
        -----
        Applies freesurfer hack for large vectors described in `issue 100`_ and
        `save_nifti.m <save77_>`_.

        Allows for freesurfer hack for 7th order icosahedron surface described
        in `issue 309`_, load_nifti.m_, and `save_nifti.m <save50_>`_.
        '''
        shape = super(Nifti1Header, self).get_data_shape()
        # Apply freesurfer hack for large vectors
        if shape[:3] == (-1, 1, 1):
            vec_len = int(self._structarr['glmin'])
            if vec_len == 0:
                raise HeaderDataError('-1 in dim[1] but 0 in glmin; '
                                      'inconsistent freesurfer type header?')
            return (vec_len, 1, 1) + shape[3:]
        # Apply freesurfer hack for ico7 surface
        elif shape[:3] == (27307, 1, 6):
            return (163842, 1, 1) + shape[3:]
        else:  # Normal case
            return shape

    def set_data_shape(self, shape):
        ''' Set shape of data  # noqa

        If ``ndims == len(shape)`` then we set zooms for dimensions higher than
        ``ndims`` to 1.0

        Nifti1 images can have up to seven dimensions. For FreeSurfer-variant
        Nifti surface files, the first dimension is assumed to correspond to
        vertices/nodes on a surface, and dimensions two and three are
        constrained to have depth of 1. Dimensions 4-7 are constrained only by
        type bounds.

        Parameters
        ----------
        shape : sequence
           sequence of integers specifying data array shape

        Notes
        -----
        Applies freesurfer hack for large vectors described in `issue 100`_ and
        `save_nifti.m <save77_>`_.

        Allows for freesurfer hack for 7th order icosahedron surface described
        in `issue 309`_, load_nifti.m_, and `save_nifti.m <save50_>`_.

        The Nifti1 `standard header`_ allows for the following "point set"
        definition of a surface, not currently implemented in nibabel.

        ::

          To signify that the vector value at each voxel is really a
          spatial coordinate (e.g., the vertices or nodes of a surface mesh):
            - dataset must have a 5th dimension
            - intent_code must be NIFTI_INTENT_POINTSET
            - dim[0] = 5
            - dim[1] = number of points
            - dim[2] = dim[3] = dim[4] = 1
            - dim[5] must be the dimensionality of space (e.g., 3 => 3D space).
            - intent_name may describe the object these points come from
              (e.g., "pial", "gray/white" , "EEG", "MEG").

        .. _issue 100: https://github.com/nipy/nibabel/issues/100
        .. _issue 309: https://github.com/nipy/nibabel/issues/309
        .. _save77:
            https://github.com/fieldtrip/fieldtrip/blob/428798b/external/freesurfer/save_nifti.m#L77-L82
        .. _save50:
            https://github.com/fieldtrip/fieldtrip/blob/428798b/external/freesurfer/save_nifti.m#L50-L56
        .. _load_nifti.m:
            https://github.com/fieldtrip/fieldtrip/blob/428798b/external/freesurfer/load_nifti.m#L86-L89
        .. _standard header: http://nifti.nimh.nih.gov/pub/dist/src/niftilib/nifti1.h
        '''
        hdr = self._structarr
        shape = tuple(shape)

        # Apply freesurfer hack for ico7 surface
        if shape[:3] == (163842, 1, 1):
            shape = (27307, 1, 6) + shape[3:]
        # Apply freesurfer hack for large vectors
        elif (len(shape) >= 3 and shape[1:3] == (1, 1) and
                shape[0] > np.iinfo(hdr['dim'].dtype.base).max):
            try:
                hdr['glmin'] = shape[0]
            except OverflowError:
                overflow = True
            else:
                overflow = hdr['glmin'] != shape[0]
            if overflow:
                raise HeaderDataError('shape[0] %s does not fit in glmax '
                                      'datatype' % shape[0])
            warnings.warn('Using large vector Freesurfer hack; header will '
                          'not be compatible with SPM or FSL', stacklevel=2)
            shape = (-1, 1, 1) + shape[3:]
        super(Nifti1Header, self).set_data_shape(shape)

    def get_qform_quaternion(self):
        ''' Compute quaternion from b, c, d of quaternion

        Fills a value by assuming this is a unit quaternion
        '''
        hdr = self._structarr
        bcd = [hdr['quatern_b'], hdr['quatern_c'], hdr['quatern_d']]
        # Adjust threshold to precision of stored values in header
        return fillpositive(bcd, self.quaternion_threshold)

    def get_qform(self, coded=False):
        """ Return 4x4 affine matrix from qform parameters in header

        Parameters
        ----------
        coded : bool, optional
            If True, return {affine or None}, and qform code.  If False, just
            return affine.  {affine or None} means, return None if qform code
            == 0, and affine otherwise.

        Returns
        -------
        affine : None or (4,4) ndarray
            If `coded` is False, always return affine reconstructed from qform
            quaternion.  If `coded` is True, return None if qform code is 0,
            else return the affine.
        code : int
            Qform code. Only returned if `coded` is True.
        """
        hdr = self._structarr
        code = int(hdr['qform_code'])
        if code == 0 and coded:
            return None, 0
        quat = self.get_qform_quaternion()
        R = quat2mat(quat)
        vox = hdr['pixdim'][1:4].copy()
        if np.any(vox < 0):
            raise HeaderDataError('pixdims[1,2,3] should be positive')
        qfac = hdr['pixdim'][0]
        if qfac not in (-1, 1):
            raise HeaderDataError('qfac (pixdim[0]) should be 1 or -1')
        vox[-1] *= qfac
        S = np.diag(vox)
        M = np.dot(R, S)
        out = np.eye(4)
        out[0:3, 0:3] = M
        out[0:3, 3] = [hdr['qoffset_x'], hdr['qoffset_y'], hdr['qoffset_z']]
        if coded:
            return out, code
        return out

    def set_qform(self, affine, code=None, strip_shears=True):
        ''' Set qform header values from 4x4 affine

        Parameters
        ----------
        affine : None or 4x4 array
            affine transform to write into sform. If None, only set code.
        code : None, string or integer, optional
            String or integer giving meaning of transform in *affine*.
            The default is None.  If code is None, then:

            * If affine is None, `code`-> 0
            * If affine not None and existing qform code in header == 0,
              `code`-> 2 (aligned)
            * If affine not None and existing qform code in header != 0,
              `code`-> existing qform code in header

        strip_shears : bool, optional
            Whether to strip shears in `affine`.  If True, shears will be
            silently stripped. If False, the presence of shears will raise a
            ``HeaderDataError``

        Notes
        -----
        The qform transform only encodes translations, rotations and
        zooms. If there are shear components to the `affine` transform, and
        `strip_shears` is True (the default), the written qform gives the
        closest approximation where the rotation matrix is orthogonal. This is
        to allow quaternion representation. The orthogonal representation
        enforces orthogonal axes.

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> int(hdr['qform_code'])  # gives 0 - unknown
        0
        >>> affine = np.diag([1,2,3,1])
        >>> np.all(hdr.get_qform() == affine)
        False
        >>> hdr.set_qform(affine)
        >>> np.all(hdr.get_qform() == affine)
        True
        >>> int(hdr['qform_code'])  # gives 2 - aligned
        2
        >>> hdr.set_qform(affine, code='talairach')
        >>> int(hdr['qform_code'])
        3
        >>> hdr.set_qform(affine, code=None)
        >>> int(hdr['qform_code'])
        3
        >>> hdr.set_qform(affine, code='scanner')
        >>> int(hdr['qform_code'])
        1
        >>> hdr.set_qform(None)
        >>> int(hdr['qform_code'])
        0
        '''
        hdr = self._structarr
        old_code = hdr['qform_code']
        if code is None:
            if affine is None:
                code = 0
            elif old_code == 0:
                code = 2  # aligned
            else:
                code = old_code
        else:  # code set
            code = self._field_recoders['qform_code'][code]
        hdr['qform_code'] = code
        if affine is None:
            return
        affine = np.asarray(affine)
        if not affine.shape == (4, 4):
            raise TypeError('Need 4x4 affine as input')
        trans = affine[:3, 3]
        RZS = affine[:3, :3]
        zooms = np.sqrt(np.sum(RZS * RZS, axis=0))
        R = RZS / zooms
        # Set qfac to make R determinant positive
        if npl.det(R) > 0:
            qfac = 1
        else:
            qfac = -1
            R[:, -1] *= -1
        # Make R orthogonal (to allow quaternion representation)
        # The orthogonal representation enforces orthogonal axes
        # (a subtle requirement of the NIFTI format qform transform)
        # Transform below is polar decomposition, returning the closest
        # orthogonal matrix PR, to input R
        P, S, Qs = npl.svd(R)
        PR = np.dot(P, Qs)
        if not strip_shears and not np.allclose(PR, R):
            raise HeaderDataError("Shears in affine and `strip_shears` is "
                                  "False")
        # Convert to quaternion
        quat = mat2quat(PR)
        # Set into header
        hdr['qoffset_x'], hdr['qoffset_y'], hdr['qoffset_z'] = trans
        hdr['pixdim'][0] = qfac
        hdr['pixdim'][1:4] = zooms
        hdr['quatern_b'], hdr['quatern_c'], hdr['quatern_d'] = quat[1:]

    def get_sform(self, coded=False):
        """ Return 4x4 affine matrix from sform parameters in header

        Parameters
        ----------
        coded : bool, optional
            If True, return {affine or None}, and sform code.  If False, just
            return affine.  {affine or None} means, return None if sform code
            == 0, and affine otherwise.

        Returns
        -------
        affine : None or (4,4) ndarray
            If `coded` is False, always return affine from sform fields. If
            `coded` is True, return None if sform code is 0, else return the
            affine.
        code : int
            Sform code. Only returned if `coded` is True.
        """
        hdr = self._structarr
        code = int(hdr['sform_code'])
        if code == 0 and coded:
            return None, 0
        out = np.eye(4)
        out[0, :] = hdr['srow_x'][:]
        out[1, :] = hdr['srow_y'][:]
        out[2, :] = hdr['srow_z'][:]
        if coded:
            return out, code
        return out

    def set_sform(self, affine, code=None):
        ''' Set sform transform from 4x4 affine

        Parameters
        ----------
        affine : None or 4x4 array
            affine transform to write into sform.  If None, only set `code`
        code : None, string or integer, optional
            String or integer giving meaning of transform in *affine*.
            The default is None.  If code is None, then:

            * If affine is None, `code`-> 0
            * If affine not None and existing sform code in header == 0,
              `code`-> 2 (aligned)
            * If affine not None and existing sform code in header != 0,
              `code`-> existing sform code in header

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> int(hdr['sform_code'])  # gives 0 - unknown
        0
        >>> affine = np.diag([1,2,3,1])
        >>> np.all(hdr.get_sform() == affine)
        False
        >>> hdr.set_sform(affine)
        >>> np.all(hdr.get_sform() == affine)
        True
        >>> int(hdr['sform_code'])  # gives 2 - aligned
        2
        >>> hdr.set_sform(affine, code='talairach')
        >>> int(hdr['sform_code'])
        3
        >>> hdr.set_sform(affine, code=None)
        >>> int(hdr['sform_code'])
        3
        >>> hdr.set_sform(affine, code='scanner')
        >>> int(hdr['sform_code'])
        1
        >>> hdr.set_sform(None)
        >>> int(hdr['sform_code'])
        0
        '''
        hdr = self._structarr
        old_code = hdr['sform_code']
        if code is None:
            if affine is None:
                code = 0
            elif old_code == 0:
                code = 2  # aligned
            else:
                code = old_code
        else:  # code set
            code = self._field_recoders['sform_code'][code]
        hdr['sform_code'] = code
        if affine is None:
            return
        affine = np.asarray(affine)
        hdr['srow_x'][:] = affine[0, :]
        hdr['srow_y'][:] = affine[1, :]
        hdr['srow_z'][:] = affine[2, :]

    def get_slope_inter(self):
        ''' Get data scaling (slope) and DC offset (intercept) from header data

        Returns
        -------
        slope : None or float
           scaling (slope).  None if there is no valid scaling from these
           fields
        inter : None or float
           offset (intercept). None if there is no valid scaling or if offset
           is not finite.

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> hdr.get_slope_inter()
        (1.0, 0.0)
        >>> hdr['scl_slope'] = 0
        >>> hdr.get_slope_inter()
        (None, None)
        >>> hdr['scl_slope'] = np.nan
        >>> hdr.get_slope_inter()
        (None, None)
        >>> hdr['scl_slope'] = 1
        >>> hdr['scl_inter'] = 1
        >>> hdr.get_slope_inter()
        (1.0, 1.0)
        >>> hdr['scl_inter'] = np.inf
        >>> hdr.get_slope_inter() #doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
            ...
        HeaderDataError: Valid slope but invalid intercept inf
        '''
        # Note that we are returning float (float64) scalefactors and
        # intercepts, although they are stored as in nifti1 as float32.
        slope = float(self['scl_slope'])
        inter = float(self['scl_inter'])
        if slope == 0 or not np.isfinite(slope):
            return None, None
        if not np.isfinite(inter):
            raise HeaderDataError(
                'Valid slope but invalid intercept {0}'.format(inter))
        return slope, inter

    def set_slope_inter(self, slope, inter=None):
        ''' Set slope and / or intercept into header

        Set slope and intercept for image data, such that, if the image
        data is ``arr``, then the scaled image data will be ``(arr *
        slope) + inter``

        (`slope`, `inter`) of (NaN, NaN) is a signal to a containing image to
        set `slope`, `inter` automatically on write.

        Parameters
        ----------
        slope : None or float
           If None, implies `slope`  of NaN. If `slope` is None or NaN then
           `inter` should be None or NaN.  Values of 0, Inf or -Inf raise
           HeaderDataError
        inter : None or float, optional
           Intercept. If None, implies `inter` of NaN. If `slope` is None or
           NaN then `inter` should be None or NaN.  Values of Inf or -Inf raise
           HeaderDataError
        '''
        if slope is None:
            slope = np.nan
        if inter is None:
            inter = np.nan
        if slope in (0, np.inf, -np.inf):
            raise HeaderDataError('Slope cannot be 0 or infinite')
        if inter in (np.inf, -np.inf):
            raise HeaderDataError('Intercept cannot be infinite')
        if np.isnan(slope) ^ np.isnan(inter):
            raise HeaderDataError('None or both of slope, inter should be nan')
        self._structarr['scl_slope'] = slope
        self._structarr['scl_inter'] = inter

    def get_dim_info(self):
        ''' Gets NIfTI MRI slice etc dimension information

        Returns
        -------
        freq : {None,0,1,2}
           Which data array axis is frequency encode direction
        phase : {None,0,1,2}
           Which data array axis is phase encode direction
        slice : {None,0,1,2}
           Which data array axis is slice encode direction

        where ``data array`` is the array returned by ``get_data``

        Because NIfTI1 files are natively Fortran indexed:
          0 is fastest changing in file
          1 is medium changing in file
          2 is slowest changing in file

        ``None`` means the axis appears not to be specified.

        Examples
        --------
        See set_dim_info function

        '''
        hdr = self._structarr
        info = int(hdr['dim_info'])
        freq = info & 3
        phase = (info >> 2) & 3
        slice = (info >> 4) & 3
        return (freq - 1 if freq else None,
                phase - 1 if phase else None,
                slice - 1 if slice else None)

    def set_dim_info(self, freq=None, phase=None, slice=None):
        ''' Sets nifti MRI slice etc dimension information

        Parameters
        ----------
        freq : {None, 0, 1, 2}
            axis of data array referring to frequency encoding
        phase : {None, 0, 1, 2}
            axis of data array referring to phase encoding
        slice : {None, 0, 1, 2}
            axis of data array referring to slice encoding

        ``None`` means the axis is not specified.

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> hdr.set_dim_info(1, 2, 0)
        >>> hdr.get_dim_info()
        (1, 2, 0)
        >>> hdr.set_dim_info(freq=1, phase=2, slice=0)
        >>> hdr.get_dim_info()
        (1, 2, 0)
        >>> hdr.set_dim_info()
        >>> hdr.get_dim_info()
        (None, None, None)
        >>> hdr.set_dim_info(freq=1, phase=None, slice=0)
        >>> hdr.get_dim_info()
        (1, None, 0)

        Notes
        -----
        This is stored in one byte in the header
        '''
        for inp in (freq, phase, slice):
            # Don't use == on None to avoid a FutureWarning in python3
            if inp is not None and inp not in (0, 1, 2):
                raise HeaderDataError('Inputs must be in [None, 0, 1, 2]')
        info = 0
        if freq is not None:
            info = info | ((freq + 1) & 3)
        if phase is not None:
            info = info | (((phase + 1) & 3) << 2)
        if slice is not None:
            info = info | (((slice + 1) & 3) << 4)
        self._structarr['dim_info'] = info

    def get_intent(self, code_repr='label'):
        ''' Get intent code, parameters and name

        Parameters
        ----------
        code_repr : string
           string giving output form of intent code representation.
           Default is 'label'; use 'code' for integer representation.

        Returns
        -------
        code : string or integer
            intent code, or string describing code
        parameters : tuple
            parameters for the intent
        name : string
            intent name

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> hdr.set_intent('t test', (10,), name='some score')
        >>> hdr.get_intent()
        ('t test', (10.0,), 'some score')
        >>> hdr.get_intent('code')
        (3, (10.0,), 'some score')
        '''
        hdr = self._structarr
        recoder = self._field_recoders['intent_code']
        code = int(hdr['intent_code'])
        known_intent = code in recoder
        if code_repr == 'code':
            label = code
        elif code_repr == 'label':
            if known_intent:
                label = recoder.label[code]
            else:
                label = 'unknown code ' + str(code)
        else:
            raise TypeError('repr can be "label" or "code"')
        n_params = len(recoder.parameters[code]) if known_intent else 0
        params = (float(hdr['intent_p%d' % (i + 1)]) for i in range(n_params))
        name = asstr(hdr['intent_name'].item())
        return label, tuple(params), name

    def set_intent(self, code, params=(), name='', allow_unknown=False):
        ''' Set the intent code, parameters and name

        If parameters are not specified, assumed to be all zero. Each
        intent code has a set number of parameters associated. If you
        specify any parameters, then it will need to be the correct number
        (e.g the "f test" intent requires 2).  However, parameters can
        also be set in the file data, so we also allow not setting any
        parameters (empty parameter tuple).

        Parameters
        ----------
        code : integer or string
            code specifying nifti intent
        params : list, tuple of scalars
            parameters relating to intent (see intent_codes)
            defaults to ().  Unspecified parameters are set to 0.0
        name : string
            intent name (description). Defaults to ''
        allow_unknown : {False, True}, optional
            Allow unknown integer intent codes. If False (the default),
            a KeyError is raised on attempts to set the intent
            to an unknown code.

        Returns
        -------
        None

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> hdr.set_intent(0)  # no intent
        >>> hdr.set_intent('z score')
        >>> hdr.get_intent()
        ('z score', (), '')
        >>> hdr.get_intent('code')
        (5, (), '')
        >>> hdr.set_intent('t test', (10,), name='some score')
        >>> hdr.get_intent()
        ('t test', (10.0,), 'some score')
        >>> hdr.set_intent('f test', (2, 10), name='another score')
        >>> hdr.get_intent()
        ('f test', (2.0, 10.0), 'another score')
        >>> hdr.set_intent('f test')
        >>> hdr.get_intent()
        ('f test', (0.0, 0.0), '')
        >>> hdr.set_intent(9999, allow_unknown=True) # unknown code
        >>> hdr.get_intent()
        ('unknown code 9999', (), '')
        '''
        hdr = self._structarr
        known_intent = code in intent_codes
        if not known_intent:
            # We can set intent via an unknown integer code, but can't via an
            # unknown string label
            if not allow_unknown or isinstance(code, string_types):
                raise KeyError('Unknown intent code: ' + str(code))
        if known_intent:
            icode = intent_codes.code[code]
            p_descr = intent_codes.parameters[code]
        else:
            icode = code
            p_descr = ('p1', 'p2', 'p3')
        if len(params) and len(params) != len(p_descr):
            raise HeaderDataError('Need params of form %s, or empty'
                                  % (p_descr,))
        hdr['intent_code'] = icode
        hdr['intent_name'] = name
        all_params = [0] * 3
        all_params[:len(params)] = params[:]
        for i, param in enumerate(all_params):
            hdr['intent_p%d' % (i + 1)] = param

    def get_slice_duration(self):
        ''' Get slice duration

        Returns
        -------
        slice_duration : float
            time to acquire one slice

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> hdr.set_dim_info(slice=2)
        >>> hdr.set_slice_duration(0.3)
        >>> print("%0.1f" % hdr.get_slice_duration())
        0.3

        Notes
        -----
        The NIfTI1 spec appears to require the slice dimension to be
        defined for slice_duration to have meaning.
        '''
        _, _, slice_dim = self.get_dim_info()
        if slice_dim is None:
            raise HeaderDataError('Slice dimension must be set '
                                  'for duration to be valid')
        return float(self._structarr['slice_duration'])

    def set_slice_duration(self, duration):
        ''' Set slice duration

        Parameters
        ----------
        duration : scalar
            time to acquire one slice

        Examples
        --------
        See ``get_slice_duration``
        '''
        _, _, slice_dim = self.get_dim_info()
        if slice_dim is None:
            raise HeaderDataError('Slice dimension must be set '
                                  'for duration to be valid')
        self._structarr['slice_duration'] = duration

    def get_n_slices(self):
        ''' Return the number of slices
        '''
        _, _, slice_dim = self.get_dim_info()
        if slice_dim is None:
            raise HeaderDataError('Slice dimension not set in header '
                                  'dim_info')
        shape = self.get_data_shape()
        try:
            slice_len = shape[slice_dim]
        except IndexError:
            raise HeaderDataError('Slice dimension index (%s) outside '
                                  'shape tuple (%s)'
                                  % (slice_dim, shape))
        return slice_len

    def get_slice_times(self):
        ''' Get slice times from slice timing information

        Returns
        -------
        slice_times : tuple
            Times of acquisition of slices, where 0 is the beginning of
            the acquisition, ordered by position in file.  nifti allows
            slices at the top and bottom of the volume to be excluded from
            the standard slice timing specification, and calls these
            "padding slices".  We give padding slices ``None`` as a time
            of acquisition

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> hdr.set_dim_info(slice=2)
        >>> hdr.set_data_shape((1, 1, 7))
        >>> hdr.set_slice_duration(0.1)
        >>> hdr['slice_code'] = slice_order_codes['sequential increasing']
        >>> slice_times = hdr.get_slice_times()
        >>> np.allclose(slice_times, [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
        True
        '''
        hdr = self._structarr
        slice_len = self.get_n_slices()
        duration = self.get_slice_duration()
        slabel = self.get_value_label('slice_code')
        if slabel == 'unknown':
            raise HeaderDataError('Cannot get slice times when '
                                  'Slice code is "unknown"')
        slice_start, slice_end = (int(hdr['slice_start']),
                                  int(hdr['slice_end']))
        if slice_start < 0:
            raise HeaderDataError('slice_start should be >= 0')
        if slice_end == 0:
            slice_end = slice_len - 1
        n_timed = slice_end - slice_start + 1
        if n_timed < 1:
            raise HeaderDataError('slice_end should be > slice_start')
        st_order = self._slice_time_order(slabel, n_timed)
        times = st_order * duration
        return ((None,) * slice_start +
                tuple(times) +
                (None,) * (slice_len - slice_end - 1))

    def set_slice_times(self, slice_times):
        ''' Set slice times into *hdr*

        Parameters
        ----------
        slice_times : tuple
            tuple of slice times, one value per slice
            tuple can include None to indicate no slice time for that slice

        Examples
        --------
        >>> hdr = Nifti1Header()
        >>> hdr.set_dim_info(slice=2)
        >>> hdr.set_data_shape([1, 1, 7])
        >>> hdr.set_slice_duration(0.1)
        >>> times = [None, 0.2, 0.4, 0.1, 0.3, 0.0, None]
        >>> hdr.set_slice_times(times)
        >>> hdr.get_value_label('slice_code')
        'alternating decreasing'
        >>> int(hdr['slice_start'])
        1
        >>> int(hdr['slice_end'])
        5
        '''
        # Check if number of slices matches header
        hdr = self._structarr
        slice_len = self.get_n_slices()
        if slice_len != len(slice_times):
            raise HeaderDataError('Number of slice times does not '
                                  'match number of slices')
        # Extract Nones at beginning and end.  Check for others
        for ind, time in enumerate(slice_times):
            if time is not None:
                slice_start = ind
                break
        else:
            raise HeaderDataError('Not all slice times can be None')
        for ind, time in enumerate(slice_times[::-1]):
            if time is not None:
                slice_end = slice_len - ind - 1
                break
        timed = slice_times[slice_start:slice_end + 1]
        for time in timed:
            if time is None:
                raise HeaderDataError('Cannot have None in middle '
                                      'of slice time vector')
        # Find slice duration, check times are compatible with single
        # duration
        tdiffs = np.diff(np.sort(timed))
        if not np.allclose(np.diff(tdiffs), 0):
            raise HeaderDataError('Slice times not compatible with '
                                  'single slice duration')
        duration = np.mean(tdiffs)
        # To slice time order
        st_order = np.round(np.array(timed) / duration)
        # Check if slice times fit known schemes
        n_timed = len(timed)
        so_recoder = self._field_recoders['slice_code']
        labels = so_recoder.value_set('label')
        labels.remove('unknown')

        matching_labels = []
        for label in labels:
            if np.all(st_order == self._slice_time_order(
                    label,
                    n_timed)):
                matching_labels.append(label)

        if not matching_labels:
            raise HeaderDataError('slice ordering of %s fits '
                                  'with no known scheme' % st_order)
        if len(matching_labels) > 1:
            warnings.warn(
                'Multiple slice orders satisfy: %s. Choosing the first one'
                % ', '.join(matching_labels)
            )
        label = matching_labels[0]
        # Set values into header
        hdr['slice_start'] = slice_start
        hdr['slice_end'] = slice_end
        hdr['slice_duration'] = duration
        hdr['slice_code'] = slice_order_codes.code[label]

    def _slice_time_order(self, slabel, n_slices):
        ''' Supporting function to give time order of slices from label '''
        if slabel == 'sequential increasing':
            sp_ind_time_order = list(range(n_slices))
        elif slabel == 'sequential decreasing':
            sp_ind_time_order = list(range(n_slices)[::-1])
        elif slabel == 'alternating increasing':
            sp_ind_time_order = (list(range(0, n_slices, 2)) +
                                 list(range(1, n_slices, 2)))
        elif slabel == 'alternating decreasing':
            sp_ind_time_order = (list(range(n_slices - 1, -1, -2)) +
                                 list(range(n_slices - 2, -1, -2)))
        elif slabel == 'alternating increasing 2':
            sp_ind_time_order = (list(range(1, n_slices, 2)) +
                                 list(range(0, n_slices, 2)))
        elif slabel == 'alternating decreasing 2':
            sp_ind_time_order = (list(range(n_slices - 2, -1, -2)) +
                                 list(range(n_slices - 1, -1, -2)))
        else:
            raise HeaderDataError('We do not handle slice ordering "%s"'
                                  % slabel)
        return np.argsort(sp_ind_time_order)

    def get_xyzt_units(self):
        xyz_code = self.structarr['xyzt_units'] % 8
        t_code = self.structarr['xyzt_units'] - xyz_code
        return (unit_codes.label[xyz_code],
                unit_codes.label[t_code])

    def set_xyzt_units(self, xyz=None, t=None):
        if xyz is None:
            xyz = 0
        if t is None:
            t = 0
        xyz_code = self.structarr['xyzt_units'] % 8
        t_code = self.structarr['xyzt_units'] - xyz_code
        xyz_code = unit_codes[xyz]
        t_code = unit_codes[t]
        self.structarr['xyzt_units'] = xyz_code + t_code

    def _clean_after_mapping(self):
        ''' Set format-specific stuff after converting header from mapping

        Clean up header after it has been initialized from an
        ``as_analyze_map`` method of another header type

        See :meth:`nibabel.analyze.AnalyzeHeader._clean_after_mapping` for a
        more detailed description.
        '''
        self._structarr['magic'] = (self.single_magic if self.is_single
                                    else self.pair_magic)

    ''' Checks only below here '''

    @classmethod
    def _get_checks(klass):
        # We need to return our own versions of - e.g. chk_datatype, to
        # pick up the Nifti datatypes from our class
        return (klass._chk_sizeof_hdr,
                klass._chk_datatype,
                klass._chk_bitpix,
                klass._chk_pixdims,
                klass._chk_qfac,
                klass._chk_magic,
                klass._chk_offset,
                klass._chk_qform_code,
                klass._chk_sform_code)

    @staticmethod
    def _chk_qfac(hdr, fix=False):
        rep = Report(HeaderDataError)
        if hdr['pixdim'][0] in (-1, 1):
            return hdr, rep
        rep.problem_level = 20
        rep.problem_msg = 'pixdim[0] (qfac) should be 1 (default) or -1'
        if fix:
            hdr['pixdim'][0] = 1
            rep.fix_msg = 'setting qfac to 1'
        return hdr, rep

    @staticmethod
    def _chk_magic(hdr, fix=False):
        rep = Report(HeaderDataError)
        magic = hdr['magic'].item()
        if magic in (hdr.pair_magic, hdr.single_magic):
            return hdr, rep
        rep.problem_msg = ('magic string "%s" is not valid' %
                           asstr(magic))
        rep.problem_level = 45
        if fix:
            rep.fix_msg = 'leaving as is, but future errors are likely'
        return hdr, rep

    @staticmethod
    def _chk_offset(hdr, fix=False):
        rep = Report(HeaderDataError)
        # for ease of later string formatting, use scalar of byte string
        magic = hdr['magic'].item()
        offset = hdr['vox_offset'].item()
        if offset == 0:
            return hdr, rep
        if magic == hdr.single_magic and offset < hdr.single_vox_offset:
            rep.problem_level = 40
            rep.problem_msg = ('vox offset %d too low for '
                               'single file nifti1' % offset)
            if fix:
                hdr['vox_offset'] = hdr.single_vox_offset
                rep.fix_msg = 'setting to minimum value of {0}'.format(
                    hdr.single_vox_offset)
            return hdr, rep
        if not offset % 16:
            return hdr, rep
        # SPM uses memory mapping to read the data, and
        # apparently this has to start on 16 byte boundaries
        rep.problem_msg = ('vox offset (={0:g}) not divisible '
                           'by 16, not SPM compatible'.format(offset))
        rep.problem_level = 30
        if fix:
            rep.fix_msg = 'leaving at current value'
        return hdr, rep

    @classmethod
    def _chk_qform_code(klass, hdr, fix=False):
        return klass._chk_xform_code('qform_code', hdr, fix)

    @classmethod
    def _chk_sform_code(klass, hdr, fix=False):
        return klass._chk_xform_code('sform_code', hdr, fix)

    @classmethod
    def _chk_xform_code(klass, code_type, hdr, fix):
        # utility method for sform and qform codes
        rep = Report(HeaderDataError)
        code = int(hdr[code_type])
        recoder = klass._field_recoders[code_type]
        if code in recoder.value_set():
            return hdr, rep
        rep.problem_level = 30
        rep.problem_msg = '%s %d not valid' % (code_type, code)
        if fix:
            hdr[code_type] = 0
            rep.fix_msg = 'setting to 0'
        return hdr, rep

    @classmethod
    def may_contain_header(klass, binaryblock):
        if len(binaryblock) < klass.sizeof_hdr:
            return False

        hdr_struct = np.ndarray(shape=(), dtype=header_dtype,
                                buffer=binaryblock[:klass.sizeof_hdr])
        return hdr_struct['magic'] in (b'ni1', b'n+1')

class Nifti1PairHeader(Nifti1Header):
    ''' Class for NIfTI1 pair header '''
    # Signal whether this is single (header + data) file
    is_single = False

class Nifti1Pair(analyze.AnalyzeImage):
    """ Class for NIfTI1 format image, header pair
    """
    header_class = Nifti1PairHeader
    _meta_sniff_len = header_class.sizeof_hdr
    rw = True

    def __init__(self, dataobj, affine, header=None,
                 extra=None, file_map=None):
        super(Nifti1Pair, self).__init__(dataobj,
                                         affine,
                                         header,
                                         extra,
                                         file_map)
        # Force set of s/q form when header is None unless affine is also None
        if header is None and affine is not None:
            self._affine2header()
    # Copy docstring
    __init__.__doc__ = analyze.AnalyzeImage.__init__.__doc__ + '''
    Notes
    -----

    If both a `header` and an `affine` are specified, and the `affine` does
    not match the affine that is in the `header`, the `affine` will be used,
    but the ``sform_code`` and ``qform_code`` fields in the header will be
    re-initialised to their default values. This is performed on the basis
    that, if you are changing the affine, you are likely to be changing the
    space to which the affine is pointing.  The :meth:`set_sform` and
    :meth:`set_qform` methods can be used to update the codes after an image
    has been created - see those methods, and the :ref:`manual
    <default-sform-qform-codes>` for more details.  '''

    def update_header(self):
        ''' Harmonize header with image data and affine

        See AnalyzeImage.update_header for more examples

        Examples
        --------
        >>> data = np.zeros((2,3,4))
        >>> affine = np.diag([1.0,2.0,3.0,1.0])
        >>> img = Nifti1Image(data, affine)
        >>> hdr = img.header
        >>> np.all(hdr.get_qform() == affine)
        True
        >>> np.all(hdr.get_sform() == affine)
        True
        '''
        super(Nifti1Pair, self).update_header()
        hdr = self._header
        hdr['magic'] = hdr.pair_magic

    def _affine2header(self):
        """ Unconditionally set affine into the header """
        hdr = self._header
        # Set affine into sform with default code
        hdr.set_sform(self._affine, code='aligned')
        # Make qform 'unknown'
        hdr.set_qform(self._affine, code='unknown')

    def get_qform(self, coded=False):
        """ Return 4x4 affine matrix from qform parameters in header

        Parameters
        ----------
        coded : bool, optional
            If True, return {affine or None}, and qform code.  If False, just
            return affine.  {affine or None} means, return None if qform code
            == 0, and affine otherwise.

        Returns
        -------
        affine : None or (4,4) ndarray
            If `coded` is False, always return affine reconstructed from qform
            quaternion.  If `coded` is True, return None if qform code is 0,
            else return the affine.
        code : int
            Qform code. Only returned if `coded` is True.

        See also
        --------
        set_qform
        get_sform
        """
        return self._header.get_qform(coded)

    def set_qform(self, affine, code=None, strip_shears=True, **kwargs):
        ''' Set qform header values from 4x4 affine

        Parameters
        ----------
        affine : None or 4x4 array
            affine transform to write into sform. If None, only set code.
        code : None, string or integer
            String or integer giving meaning of transform in *affine*.
            The default is None.  If code is None, then:

            * If affine is None, `code`-> 0
            * If affine not None and existing qform code in header == 0,
              `code`-> 2 (aligned)
            * If affine not None and existing qform code in header != 0,
              `code`-> existing qform code in header

        strip_shears : bool, optional
            Whether to strip shears in `affine`.  If True, shears will be
            silently stripped. If False, the presence of shears will raise a
            ``HeaderDataError``
        update_affine : bool, optional
            Whether to update the image affine from the header best affine
            after setting the qform. Must be keyword argument (because of
            different position in `set_qform`). Default is True

        See also
        --------
        get_qform
        set_sform

        Examples
        --------
        >>> data = np.arange(24).reshape((2,3,4))
        >>> aff = np.diag([2, 3, 4, 1])
        >>> img = Nifti1Pair(data, aff)
        >>> img.get_qform()
        array([[ 2.,  0.,  0.,  0.],
               [ 0.,  3.,  0.,  0.],
               [ 0.,  0.,  4.,  0.],
               [ 0.,  0.,  0.,  1.]])
        >>> img.get_qform(coded=True)
        (None, 0)
        >>> aff2 = np.diag([3, 4, 5, 1])
        >>> img.set_qform(aff2, 'talairach')
        >>> qaff, code = img.get_qform(coded=True)
        >>> np.all(qaff == aff2)
        True
        >>> int(code)
        3
        '''
        update_affine = kwargs.pop('update_affine', True)
        if kwargs:
            raise TypeError('Unexpected keyword argument(s) %s' % kwargs)
        self._header.set_qform(affine, code, strip_shears)
        if update_affine:
            if self._affine is None:
                self._affine = self._header.get_best_affine()
            else:
                self._affine[:] = self._header.get_best_affine()

    def get_sform(self, coded=False):
        """ Return 4x4 affine matrix from sform parameters in header

        Parameters
        ----------
        coded : bool, optional
            If True, return {affine or None}, and sform code.  If False, just
            return affine.  {affine or None} means, return None if sform code
            == 0, and affine otherwise.

        Returns
        -------
        affine : None or (4,4) ndarray
            If `coded` is False, always return affine from sform fields. If
            `coded` is True, return None if sform code is 0, else return the
            affine.
        code : int
            Sform code. Only returned if `coded` is True.

        See also
        --------
        set_sform
        get_qform
        """
        return self._header.get_sform(coded)

    def set_sform(self, affine, code=None, **kwargs):
        ''' Set sform transform from 4x4 affine

        Parameters
        ----------
        affine : None or 4x4 array
            affine transform to write into sform.  If None, only set `code`
        code : None, string or integer
            String or integer giving meaning of transform in *affine*.
            The default is None.  If code is None, then:

            * If affine is None, `code`-> 0
            * If affine not None and existing sform code in header == 0,
              `code`-> 2 (aligned)
            * If affine not None and existing sform code in header != 0,
              `code`-> existing sform code in header

        update_affine : bool, optional
            Whether to update the image affine from the header best affine
            after setting the qform.  Must be keyword argument (because of
            different position in `set_qform`). Default is True

        See also
        --------
        get_sform
        set_qform

        Examples
        --------
        >>> data = np.arange(24).reshape((2,3,4))
        >>> aff = np.diag([2, 3, 4, 1])
        >>> img = Nifti1Pair(data, aff)
        >>> img.get_sform()
        array([[ 2.,  0.,  0.,  0.],
               [ 0.,  3.,  0.,  0.],
               [ 0.,  0.,  4.,  0.],
               [ 0.,  0.,  0.,  1.]])
        >>> saff, code = img.get_sform(coded=True)
        >>> saff
        array([[ 2.,  0.,  0.,  0.],
               [ 0.,  3.,  0.,  0.],
               [ 0.,  0.,  4.,  0.],
               [ 0.,  0.,  0.,  1.]])
        >>> int(code)
        2
        >>> aff2 = np.diag([3, 4, 5, 1])
        >>> img.set_sform(aff2, 'talairach')
        >>> saff, code = img.get_sform(coded=True)
        >>> np.all(saff == aff2)
        True
        >>> int(code)
        3
        '''
        update_affine = kwargs.pop('update_affine', True)
        if kwargs:
            raise TypeError('Unexpected keyword argument(s) %s' % kwargs)
        self._header.set_sform(affine, code)
        if update_affine:
            if self._affine is None:
                self._affine = self._header.get_best_affine()
            else:
                self._affine[:] = self._header.get_best_affine()

    def as_reoriented(self, ornt):
        """Apply an orientation change and return a new image

        If ornt is identity transform, return the original image, unchanged

        Parameters
        ----------
        ornt : (n,2) orientation array
           orientation transform. ``ornt[N,1]` is flip of axis N of the
           array implied by `shape`, where 1 means no flip and -1 means
           flip.  For example, if ``N==0`` and ``ornt[0,1] == -1``, and
           there's an array ``arr`` of shape `shape`, the flip would
           correspond to the effect of ``np.flipud(arr)``.  ``ornt[:,0]`` is
           the transpose that needs to be done to the implied array, as in
           ``arr.transpose(ornt[:,0])``
        """
        img = super(Nifti1Pair, self).as_reoriented(ornt)

        if img is self:
            return img

        # Also apply the transform to the dim_info fields
        new_dim = list(img.header.get_dim_info())
        for idx, value in enumerate(new_dim):
            # For each value, leave as None if it was that way,
            # otherwise check where we have mapped it to
            if value is None:
                continue
            new_dim[idx] = np.where(ornt[:, 0] == idx)[0]

        img.header.set_dim_info(*new_dim)

        return img

class Nifti1Image(Nifti1Pair):
    """ Class for single file NIfTI1 format image
    """
    header_class = Nifti1Header
    valid_exts = ('.nii',)
    files_types = (('image', '.nii'),)

    @staticmethod
    def _get_fileholders(file_map):
        """ Return fileholder for header and image

        For single-file niftis, the fileholder for the header and the image
        will be the same
        """
        return file_map['image'], file_map['image']

    def update_header(self):
        ''' Harmonize header with image data and affine '''
        super(Nifti1Image, self).update_header()
        hdr = self._header
        hdr['magic'] = hdr.single_magic

def load(filename):
    """ Load NIfTI1 single or pair from `filename`

    Parameters
    ----------
    filename : str
        filename of image to be loaded

    Returns
    -------
    img : Nifti1Image or Nifti1Pair
        NIfTI1 single or pair image instance

    Raises
    ------
    ImageFileError
        if `filename` doesn't look like NIfTI1;
    IOError
        if `filename` does not exist.
    """
    try:
        img = Nifti1Image.load(filename)
    except ImageFileError:
        return Nifti1Pair.load(filename)
    return img

def save(img, filename):
    """ Save NIfTI1 single or pair to `filename`

    Parameters
    ----------
    filename : str
        filename to which to save image
    """
    try:
        Nifti1Image.instance_to_filename(img, filename)
    except ImageFileError:
        Nifti1Pair.instance_to_filename(img, filename)

max_stars_repo_name: Raniac/NEURO-LEARN

max_forks_repo_name: Raniac/NEURO-LEARN

reason: import numpy


text:
import numpy as np

class Body:
	def __init__( self, x0 = 0.0, y0 = 0.0, radius = 2.0, theta0 = 0, dtheta0 = 0):
		v0 = 0.0

		self.theta0 = theta0
		self.dtheta0 = dtheta0
		self.v0 = v0
		self.x0 = np.array([x0, y0])

		self.dtheta = dtheta0
		self.radius = radius
		self.v = v0
		self.dx = v0

		self.gamma = 100.0
		self.mu = 2.0
		self.m = 1.0
		self.kappa = 0.08
		self.A = 2*np.pi*radius

		self.x = np.array([x0, y0])
		self.theta = theta0

	def reset( self ):
		self.theta = self.theta0
		self.dtheta = self.dtheta0
		self.v = self.v0
		self.x = self.x0

	def step( self, h, drive_wheels ):
	# drive_wheels expected to be a two dimensional vector v \in [0,1]x[0,1]
		print drive_wheels
		d = self.radius
		Ic = self.m*self.kappa**2.0
		F1 = self.gamma*drive_wheels[0,0]
		F2 = self.gamma*drive_wheels[0,1]

		Fmu = self.mu*np.linalg.norm(self.v) if \
			  (np.linalg.norm(self.v) > 0.0)  else 0.0
		#and (F1 < 0.1 and F2 < 0.1)

		Fmu_theta = 10*self.dtheta if \
			  (np.linalg.norm(self.dtheta) > 0.0) else 0.0

		F_left = F1 - Fmu
		F_right = F2 - Fmu
		dv = (F_left + F_right)/self.m
		ddtheta = (d*F_right - d*F_left)/Ic - Fmu_theta

		# if all(ddtheta) > 0.0:
		# 	r = -dv/ddtheta
		# else:
		# 	r = x

		self.v += h*dv
		self.dtheta += h*ddtheta
		self.theta += h*self.dtheta
		self.dx = np.array([self.v*np.cos(self.theta),
			  				self.v*np.sin(self.theta)])

		self.x += h*self.dx
		return self.x, self.dx # Returns current position

	def getPosition( self ):
	# Returns the postion as a tuple
		return self.x[0], self.x[1]

	def getOrientation( self ):
		return self.theta

	def clone( self ):
		b = Body()
		b.x =  np.array([self.x[0], self.x[1]])
		b.theta = self.theta
		b.radius = self.radius

		return b

ext: py

max_stars_repo_name: ABRG-Models/MammalBot

max_issues_repo_name: ABRG-Models/MammalBot

max_forks_repo_name: ABRG-Models/MammalBot


text:
import numpy as np
import os
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from scipy.integrate import solve_ivp

class DoublePendulum():
    """Class for modeling a double pendulum."""
    def __init__(self, L1=1, L2=1, M1=1, M2=1, g=9.81):
        """
        The class should be called with the following arguments:
        Paramenters: M1 -- Mass of inner pendulum (closest to center) [kg]
                     M2 -- Mass of outer pendulum (furthest from center) [kg]
                     L1 -- Length of massless rod between center and M1 [m]
                     L2 -- Length of massless rod between M1 and M2 [m]
                     g -- Gravitational acceleration [m/s^2]
        """
        self.L1 = L1
        self.L2 = L2
        self.M1 = M1
        self.M2 = M2

        self.g = g

    def __call__(self, t, y):
        """
        __call__ special method for caluculating the derivative of the angles
        theta1 and theta2 [angle(unitless)] and the angular velocities omega1
        and omega2 [1/s] of the dobble pendulum.
        Paramenters: t -- Time [s]
                     y -- Tuple consisting of initial values for theta1 [],
                          theta2 [], omega1 [1/s] and omega2 [1/s]. Should be
                          given as a tuple (theta1, theta2, omega1, omega2).
        Returns: Derivative of angles and angular velocities.
        """
        L1 = self.L1
        L2 = self.L2
        M1 = self.M1
        M2 = self.M2
        g  = self.g

        theta1, theta2, omega1, omega2 = y # y[0], y[1], y[2], y[3]
        diff_theta = theta2 - theta1       # Difference in angles

        topp1 = (M2*L1*omega1**2*np.sin(diff_theta)*np.cos(diff_theta)
                 + M2*g*np.sin(theta2)*np.cos(diff_theta)
                 + M2*L2*omega2**2*np.sin(diff_theta)
                 - (M1 + M2)*g*np.sin(theta1))

        bunn1 = (M1 + M2)*L1 - M2*L1*np.cos(diff_theta)**2

        topp2 = (-M2*L2*omega2**2*np.sin(diff_theta)*np.cos(diff_theta)
                 + (M1 + M2)*g*np.sin(theta1)*np.cos(diff_theta)
                 - (M1 + M2)*L1*omega1**2*np.sin(diff_theta)
                 - (M1 + M2)*g*np.sin(theta2))

        bunn2 = (M1 + M2)*L2 - M2*L2*np.cos(diff_theta)**2

        d_theta1__dt = omega1      # Angular velocity omega
        d_theta2__dt = omega2      # Angular velocity omega
        d_omega1__dt = topp1/bunn1 # Angular acceleration
        d_omega2__dt = topp2/bunn2 # Angular acceleration

        return d_theta1__dt, d_theta2__dt, d_omega1__dt, d_omega2__dt

    def solve(self, y0, T, dt, angles="rad"):
        """
        Solves the ODE using scipy.integrate solve_ivp function.
        Parameters: y0 -- An initial values for the function y we want to find
                          at a time t=0. Should be given as a list
                          [theta1, theta2, omega1, omega2].
                    T  -- End time for the experiment
                    dt -- Timestep we want to use in our calculation.
                    (NB! Both T and dt are given as scalar values. That is,
                    solve method should be called as
                    object.solve([theta1, theta2, omega1, omega2], T, dt), where
                    theta1, theta2, omega1, omega2, T and dt are scalar values.)
        Kwargs: angles -- specifying if the given values of theta and omega
                          are given in degrees or radians. Should be specified
                          as "deg" or "rad". If angles="deg", it will be changed
                          to radians and the result will allways be given in
                          radians.
        Returns: Nothing. Sets the solution as a private variable _solution.
        """
        self._dt = dt
        self._T = T

        if angles == "deg":
            y0 = np.array(y0)
            y0 = y0*(np.pi/180) # Changing from degrees to radians.

        sol = solve_ivp(self, [0, T], y0, t_eval=np.arange(0, T, dt),method="Radau")
        #store solutions
        self._solution = sol

    def create_animation(self,axis=[-3,3,-3,3]):
        """
        Creates animation of the pendulums, and traces the movement of the
        lower pendulum.
        Parameters: No parameters
        Kwargs:     axis -- A list containing minimum and maximum x and y values
                            for the coordinate system.
        Returns: Nothing. Creates an animation that can be:
                                - saved with save_animation()
                                - showed with show_animation()
        """
        fig = plt.figure()
        ax = fig.add_subplot(111)

        #make it pretty
        plt.axis("equal")
        plt.axis("off")
        plt.axis(axis)

        # Make an 'empty' plot object to be updated throughout the animation
        self.pendulums, = ax.plot([],[],'ro-',lw=2)

        #trace movement of lower pendulum
        self.trace, = ax.plot([],[],'b-')
        self.foox = [] ; self.fooy = []

        #timer
        self.time_template = 'time = %.1f s'
        self.time_text = ax.text(0.05,0.9,'',transform=ax.transAxes)

        #animate
        self.animation = animation.FuncAnimation(fig,self._next_frame,
                                                 frames=range(0,len(self.t),2),
                                                 repeat=None,
                                                 interval=1000*self.dt,
                                                 blit=True)

    def _next_frame(self,i):
        """
        Method used by create_animation() to update the plot for each frame.
        Parametes:  i -- the current iteration/frame, automatically updated
                         by the FuncAnimation() method.
        Returns: the updated plots and the new time text.
        """
        self.pendulums.set_data((0,self.x1[i],self.x2[i]),
                                (0,self.y1[i],self.y2[i]))

        self.time_text.set_text(self.time_template % (i*dt))
        self.foox.append(self.x2[i]) ; self.fooy.append(self.y2[i])
        self.trace.set_data(self.foox,self.fooy)

        return self.pendulums,self.time_text,self.trace

    def show_animation(self):
        """Shows the animation"""
        plt.show()

    def save_animation(self,filename="pendel.mp4",res=100):
        """
        Saves the animation as desired filename
        Parameters: (None)
        Kwargs:    filename -- String containing desired filename: "example.mp4"
                               Set to "pendel.mp4" as default.
                   res      -- Set default to dpi = 100, changes resolution.
        Returns: Nothing. Saves the file as specified in filename.
        """
        if os.path.isfile(filename):
            os.system("rm "+filename)
        self.animation.save(filename,fps=60,dpi=res)

    @property
    def solution(self):
        """Returns the latest stored solution."""
        try:
            return self._solution
        except AttributeError:
            raise AttributeError("Solutions for the double pendulum does not "
                                 "exist. You need to run the solve method to "
                                 "compute how the pendulum moves. Run solve "
                                 "method first.")

    @property
    def dt(self):
        """Returns dt as given in solve."""
        return self._dt

    @property
    def T(self):
        """Returns T as given in solve."""
        return self._T

    @property
    def t(self):
        """
        Returns the time calculated in solve method. Raises error if it doesn't
        exit.
        """
        try:
            return self.solution.t
        except AttributeError:
            raise AttributeError("Pendulum time does not exist. You need to "
                                  "run the solve method to compute a time. Run "
                                  "solve method first.")

    @property
    def theta1(self):
        """Returns theta of upper pendulum, raises error if it doesn't exist."""
        try:
            return self.solution.y[0]
        except AttributeError:
            raise AttributeError("Theta1 values for pendulum does not exist "
                                  "yet. You need to run the solve method to "
                                  "compute the theta1 values. Run solve method "
                                  "first.")

    @property
    def theta2(self):
        """Returns theta of lower pendulum, raises error if it doesn't exist."""
        try:
            return self.solution.y[1]
        except AttributeError:
            raise AttributeError("Theta2 values for pendulum does not exist "
                                 "yet. You need to run the solve method to "
                                 "compute the theta2 values. Run solve method "
                                 "first.")

    @property
    def omega1(self):
        """Returns omega of upper pendulum, raises error if it doesn't exist."""
        try:
            return self.solution.y[2]
        except AttributeError:
            raise AttributeError("Omega1 values for pendulum does not exist "
                                 "yet. You need to run the solve method to "
                                 "compute the omega1 values. Run solve method "
                                 "first.")

    @property
    def omega2(self):
        """Returns omega of lower pendulum, raises error if it doesn't exist."""
        try:
            return self.solution.y[3]
        except AttributeError:
            raise AttributeError("Omega1 values for pendulum does not exist "
                                 "yet. You need to run the solve method to "
                                 "compute the omega2 values. Run solve method "
                                 "first.")

    @property
    def x1(self):
        """Returns x position of upper pendulum"""
        theta1 = self.theta1
        x1_pos = self.L1*np.sin(theta1)
        return x1_pos

    @property
    def x2(self):
        """Returns x position of lower pendulum"""
        theta2 = self.theta2
        x2_pos = self.x1 + self.L2*np.sin(theta2)
        return x2_pos

    @property
    def y1(self):
        """Returns y position of upper pendulum"""
        theta1 = self.theta1
        y1_pos = -self.L1*np.cos(theta1)
        return y1_pos

    @property
    def y2(self):
        """Returns y position of lower pendulum"""
        theta2 = self.theta2
        y2_pos = self.y1 - self.L2*np.cos(theta2)
        return y2_pos

    @property
    def vx1(self):
        """Retuns speed in x direction for upper pendulum"""
        return np.gradient(self.x1,self.t)

    @property
    def vx2(self):
        """Retuns speed in x direction for lower pendulum"""
        return np.gradient(self.x2,self.t)

    @property
    def vy1(self):
        """Retuns speed in y direction for upper pendulum"""
        return np.gradient(self.y1,self.t)

    @property
    def vy2(self):
        """Retuns speed in y direction for lower pendulum"""
        return np.gradient(self.y2,self.t)

    @property
    def potential(self):
        """Returns calculated total potential energy """
        P1 = self.M1*self.g*(self.y1 + self.L1)
        P2 = self.M2*self.g*(self.y2 + self.L1 + self.L2)
        P = P1 + P2
        return P

    @property
    def kinetic(self):
        """Returns calculated total kinetic energy """
        K1 = 0.5*self.M1*(self.vx1**2 + self.vy1**2)
        K2 = 0.5*self.M2*(self.vx2**2 + self.vy2**2)
        K = K1 + K2
        return K

if __name__ == "__main__":
    #Plotting double penulum, movement and energy levels
    y0 = [-np.pi/2,-np.pi,-0.6, 0]
    T  = 10
    dt = float(1/120)

    pendel = DoublePendulum()
    pendel.solve(y0,T,dt)

    time = pendel.t
    Ke   = pendel.kinetic
    Pe   = pendel.potential
    Te   = Ke + Pe

    #plot Energy and time
    plt.style.use("classic")
    plt.plot(time,Ke)
    plt.plot(time,Pe)
    plt.plot(time,Te)
    plt.title("Double pendulum")
    plt.legend(["Kinetic energy","Potential energy","Total energy"])
    plt.xlabel("time [s]") ; plt.ylabel("Energy [J]")
    plt.grid()
    plt.show()

    #create, save and show animation
    pendel.create_animation()
    #pendel.save_animation("pendulum.mp4", res=200)
    pendel.show_animation()

reason: import numpy,from scipy


text:
import sys, os, math
import numpy as np
import pandas as pd
import scipy, scipy.fft
import scipy.interpolate

import mrc
import pickle
import json

import torch
import torch.nn.functional as F

import matplotlib
import matplotlib.pyplot as plt

from tqdm import tqdm, trange

from focustools.util import CorrectCTF, Project, BackProject

def to_pytorch_complex(x):
    return torch.stack([torch.Tensor(x.real), torch.Tensor(x.imag)], 0) # [2, ...]

def to_numpy_complex(x):
    x = x.cpu().numpy()
    return x[0] + 1j * x[1]

def pytorch_complex_mul(x1, x2):
    return torch.stack([x1[0]*x2[0]-x1[1]*x2[1], x1[0]*x2[1]+x1[1]*x2[0]], 0) # [2, ...]

def pytorch_complex_expi(x):
    return torch.stack([torch.cos(x), torch.sin(x)], 0) # [2, ...]

def pytorch_complex_conj(x):
    return torch.stack([x[0], -x[1]], 0) # [2, ...]

def pytorch_complex_abs(x):
    return torch.sqrt(x[0]**2 + x[1]**2) # [...]

def pytorch_fftfreq(n):
    if n % 2 == 0:
        raster = torch.linspace(-n//2, n//2-1, n)
    else:
        raster = torch.linspace(-(n-1)//2, (n-1)//2, n)
    return raster / n

def pytorch_phase_shift(Fx, shift):
    grids = torch.stack(torch.meshgrid([pytorch_fftfreq(n) for n in Fx.shape[1:]]), -1)

    # exp{-2j*pi*(k*x)/N}
    phase_s = torch.sum(grids * torch.Tensor(shift), -1)
    phase_s = -2 * np.pi * phase_s
    phase_s = pytorch_complex_expi(phase_s) # [2, ...]

    # complex multiplication
    return pytorch_complex_mul(Fx, phase_s) # [2, ...]

def pytorch_affine_transform(vol, theta, interp_mode='bilinear'):

    # transpose for align axis with [y, x, z] (the indexing order)
    vol = vol.permute([0, 2, 3, 1])

    if theta.shape[-1] == 3: # 3 x 3 rotation matrix
        theta = torch.cat([theta, torch.zeros_like(theta[..., :1])], -1) # (3, 4)

    input = vol[None, ...] # [N=1, C=2, D, H, W]
    grids = F.affine_grid(theta[None, ...], input.shape, align_corners=False)
    output = F.grid_sample(input, grids, mode=interp_mode, padding_mode='zeros', align_corners=False) # [N=1, C=2, D, H, W]
    output = output.squeeze(0) # [C=2, D, H, W]

    return output.permute([0, 3, 1, 2])

def volume_project(vol, pose, ffted=False, plotax=None):

    # fast Fourier transform
    if not ffted:
        # twice size to prevent aliasing
        pad_shape = [math.ceil(s * 2) for s in vol.shape]
        shift_shape = np.array([-math.ceil(s // 2) for s in vol.shape], dtype=np.int32)

        Fvol = scipy.fft.fftn(vol, s=pad_shape, workers=-1) # all cpu
        Fvol = scipy.fft.fftshift(Fvol)
    else:
        Fvol = vol
        shift_shape = None

    Fvol = to_pytorch_complex(Fvol)
    if shift_shape is not None:
        Fvol = pytorch_phase_shift(Fvol, shift_shape) # [2, D, H, W]

    # R = torch.Tensor(euler_to_rotmat(pose[0], pose[1], pose[2])[0])
    R = torch.Tensor(pose_to_rotmat(pose[0], pose[1], pose[2])[0])
    Fvol = pytorch_affine_transform(Fvol, R, interp_mode='nearest') # [2, D, H, W]

    # slice orthogonal to x-axis
    slice_pos = Fvol.shape[1]//2
    Fslice = Fvol[:, slice_pos, :, :]

    # To accelerate, we can first slice then perform shifting back
    # because the slice is aligned with the y-z plane
    if shift_shape is not None:
        Fslice = pytorch_phase_shift(Fslice, -shift_shape[1:]) # [2, D, H, W]

    # optional: translation using phase shift
    if pose[3] != 0 and pose[4] != 0:
        Fslice = pytorch_phase_shift(Fslice, [pose[3], pose[4]])

    if plotax is not None:
        pm = plotax.imshow(pytorch_complex_abs(Fslice).cpu().numpy())
        plt.colorbar(pm, ax=plotax)

    Fslice = to_numpy_complex(Fslice)
    Fslice = scipy.fft.ifftshift(Fslice)
    slice = scipy.fft.ifftn(Fslice)

    # remember to crop back
    return slice.real[:vol.shape[1], :vol.shape[2]]

def volume_backproject(img, pose, ffted=False, return_fft=False, return_weights=False, plotax=None):

    # fast Fourier transform
    if not ffted:
        # twice size to prevent aliasing
        pad_shape = [math.ceil(s * 2) for s in img.shape]
        shift_img = [math.ceil(s // 2) for s in img.shape]

        Fimg = scipy.fft.fftn(img, s=pad_shape, workers=-1) # all cpu
        Fimg = scipy.fft.fftshift(Fimg)
    else:
        Fimg = img
        # shift_img = None

    Fimg = to_pytorch_complex(Fimg)

    if plotax is not None:
        pm = plotax.imshow(pytorch_complex_abs(Fimg).cpu().numpy())
        plt.colorbar(pm, ax=plotax)

    # optional: translation using phase shift
    if pose[3] != 0 and pose[4] != 0:
        Fimg = pytorch_phase_shift(Fimg, [-pose[3], -pose[4]])

    # shift image to center
    if shift_img is not None:
        Fimg = pytorch_phase_shift(Fimg, shift_img)

    # construct empty volume (suppose square D = H)
    Fvol = torch.zeros(3, Fimg.shape[1], Fimg.shape[1], Fimg.shape[2]) # [real+imag+mask, D, H, W]

    # put back to the y-z plane
    slice_pos = Fvol.shape[1]//2
    Fvol[0, slice_pos, ...] = Fimg[0] # real
    Fvol[1, slice_pos, ...] = Fimg[1] # imag
    Fvol[2, slice_pos, ...] = 1.      # mask

    # R = torch.Tensor(euler_to_rotmat(pose[0], pose[1], pose[2])[0].T) # inverse
    R = torch.Tensor(pose_to_rotmat(pose[0], pose[1], pose[2])[0].T)
    Fvol = pytorch_affine_transform(Fvol, R, interp_mode='nearest') # [3, D, H, W]
    Fvol, W = Fvol[:2, ...], Fvol[2, ...] # [2, D, H, W], [D, H, W]

    if return_fft:
        Fvol = to_numpy_complex(Fvol)
        if return_weights:
            return Fvol, W.cpu().numpy()
        else:
            return Fvol
    else:
        # shift back
        vol_shape = (img.shape[0], img.shape[0], img.shape[1])
        if shift_img is not None:
            shift_vol = [0]+[math.ceil(s // 2) for s in vol_shape[1:]]
            Fvol = pytorch_phase_shift(Fvol, shift_vol) # [2, D, H, W]

        Fvol = to_numpy_complex(Fvol)
        Fvol = scipy.fft.ifftshift(Fvol)
        Fvol = scipy.fft.ifftn(Fvol)
        return Fvol.real[:vol_shape[0], :vol_shape[1], :vol.shape[2]]

# Sanity check
if __name__ == "__main__":
    fig = plt.figure(figsize=(10, 3))
    ax0, ax1, ax2 = fig.subplots(1, 3)

    ax0.imshow(img, cmap='gray')
    ax0.axis('off')

    pad_shape = [s * 2 for s in img.shape]
    shift_shape = np.array([-math.ceil(s // 2) for s in img.shape], dtype=np.int32)
    Fimg = scipy.fft.fftn(img, s=pad_shape, workers=-1)
    Fimg = scipy.fft.fftshift(Fimg)
    # Fimg = to_pytorch_complex(Fimg)
    # Fimg = pytorch_phase_shift(Fimg, shift_shape)

    Fvol = np.zeros((Fimg.shape[0], Fimg.shape[0], Fimg.shape[1]), dtype=np.complex64)
    Fvol[Fvol.shape[0]//2, :, :] = Fimg

    Fslice = Fvol[Fvol.shape[0]//2, :, :]
    # Fslice = pytorch_phase_shift(Fimg, -shift_shape)
    Fslice = scipy.fft.ifftshift(Fslice)
    slice = scipy.fft.ifftn(Fslice).real
    ax1.imshow(slice[:img.shape[0], :img.shape[1]], cmap='gray')
    ax1.axis('off')

    R = torch.Tensor(euler_to_rotmat(75.676, 8.5, 230.5)[0])
    Fvol = to_pytorch_complex(Fvol)
    Fvol = pytorch_affine_transform(Fvol, R.T) # [2, D, H, W]
    Fvol = pytorch_affine_transform(Fvol, R) # [2, D, H, W]
    Fvol = to_numpy_complex(Fvol)

    Fslice = Fvol[Fvol.shape[0]//2, :, :]
    # Fslice = pytorch_phase_shift(Fimg, -shift_shape)
    Fslice = scipy.fft.ifftshift(Fslice)
    slice = scipy.fft.ifftn(Fslice).real
    ax2.imshow(slice[:img.shape[0], :img.shape[1]], cmap='gray')
    ax2.axis('off')

max_stars_repo_name: peihaowang/nerf-pytorch

reason: import numpy,import scipy