Source code for pyFC.clouds

"""
.. module:: clouds
   :synopsis: Sub-module containing all the main classes for creating and transforming fractal cubes.


.. moduleauthor:: Alexander Y. Wagner <alexaner.y.wagner@gmail.com>


To do:

  - Ensure that the routine works in 2D and 1D (Allow n{ijk}=1)
  - Parallelize (with mpi4py?).

"""
import os.path as osp
import copy
from collections import OrderedDict

import numpy as np
import numpy.fft as fft

import pyFC.parsetools as pt
import pyFC.mathtools as mt

import scipy.stats as sps
import scipy.special as spf


# TODO Ensure that the routine works in 2D and 1D (Allow n{ijk}=1)
# TODO Parallelize (with mpi4py?).


class FCSlicer():
    """
    Slices fractal cube object.
    """

    def __init__(self):
        pass

    def slice(self, fc, ax=0, loc=0.5, scale='frac'):
        """
        A simple mid-plane slice.

        :arg obj fc:        :class:`pyFC.FractalCube` object
        :arg int ax:        Axis number perpendicular to slice plane
        :arg int loc:       Location of slice on that axis
        :arg str scale:     {'idx'|'frac'} integer index (idx) or fraction of cube width (rounded to nearest integer,
                            frac) for location. Default 'frac'.

        :return:            2D array of slice
        :rtype:             ndarray
        """

        if scale == 'frac': loc = int(np.round(loc * fc.cube.shape[ax]))

        return np.rollaxis(fc.cube, ax)[loc, :, :]

    def tri_slice(self, fc, locs=(0.5, 0.5, 0.5), scale='frac'):
        """ 
        A generator function for three perpendicular slices at point loc

        :arg obj fc:        :class:`pyFC.FractalCube` object
        :arg tup locs:      3-Tuple indicating location of tri-slice (point where planes intersect)
        :arg str scale:     {'idx'|'frac'} integer index (idx) or fraction of cube width (rounded to nearest integer,
                            frac) for location. Default 'frac'.

        :return:            Yield slice arrays for each axis and intercept. This will always yield three
                            slices, regardless of whether cube is 2D
        :rtype:             ndarray
        """

        if scale == 'frac': locs = np.round(locs * np.array(fc.cube.shape))

        for i, loc in enumerate(locs):
            yield np.rollaxis(fc.cube, i)[loc, :, :]


class FCAffine():
    """
    Affine transformations on a fractal cube
    """

    def __init__(self):
        pass

    def translate(self, fc, ax=0, delta=0.5, scale='frac', out='copy'):
        """ 
        Translation of cube using np.roll.

        :arg obj fc:        :class:`pyFC.FractalCube` object.
        :arg int ax:        Axis number along which to translate.
        :arg flt delta:     Distance of translation.
        :arg flt scale:     {'idx'|'frac'} integer index (idx) or fraction of cube width (rounded to
                            nearest integer, frac) for location. Default 'frac'.
        :arg str out:       One of the modes accepted by fc._returner.

        :return:            Translated cube
        :rtype:             Type depends on 'out'. By default it returns a copy
                            of the :class:`pyFC.FractalCube` object.
        
        """

        if scale == 'frac': delta = int(np.round(delta * fc.cube.shape[ax]))

        result = np.roll(fc.cube, delta, axis=ax)
        return fc._returner(result, out)

    def permute(self, fc, choice=0, out='copy'):
        """ 
        Chooses one of six permutations of the cube's axis orders.

        :arg obj fc:        :class:`pyFC.FractalCube` object.
        :arg int choice:    Choses one of six permuations::

                                0: 012, 1: 210 (= 0.T), 
                                2: 120, 3: 021 (= 2.T), 
                                4: 201, 5: 102 (= 4.T), 

                            where T is the transpose.

        :arg str out:       One of the modes accepted by fc._returner.

        :return:            Permuted cube.
        :rtype:             Type depends on 'out'. By default it returns a copy
                            of the :class:`pyFC.FractalCube` object.
 
        """

        if choice % 2 == 0:
            result = np.rollaxis(fc.cube, choice % 3)
        else:
            result = np.rollaxis(fc.cube, choice % 3).T

        return fc._returner(result, out)

    def mirror(self, fc, ax=0, out='copy'):
        """
        Mirrors a cube about a face along axis ax.

        :arg obj fc:        :class:`pyFC.FractalCube` object.
        :arg int ax:        Axis number perpendicular to which the mirror plane lies.
        :arg str out:       The "out mode". One of the modes accepted by fc._returner. Possible values are
                            'inplace':  save the ndarray data to self.cube return the FractalCube instance;
                            'ndarray':  save the ndarray data to self.cube and return the ndarray;
                            'copy':     return a copy of the FractalCube instance, with the ndarray data saved to the copy's self.cube.
        :return:            Mirrored cube
        :rtype:             Type depends on 'out'. By default it returns a copy
                            of the :class:`pyFC.FractalCube` object.

        """
        if ax == 0: result = fc.cube[::-1, :, :]
        if ax == 1: result = fc.cube[:, ::-1, :]
        if ax == 2: result = fc.cube[:, :, ::-1]

        return fc._returner(result, out)

    def center(self, fc, ax=None, mode='max', point=None, out='copy', return_point=False):
        """
        Centers a cube at a specific point. The translation is done with the :func:pyFC.FCAffine.translate.
        Weights are logarithmic.

        :arg obj fc:        :class:`pyFC.FractalCube` object.
        :arg int ax:        Axis number along which to center. If None, it will center
                            along all axes
        :arg str mode:      One of 'max', 'av_max', 'average', 'point'
                            Sets how the center point is determined:

                                -'max':         maximum value along an axis (or in entire cube if ax==None)
                                -'mean_max':     the mean location of all maxima along one direction
                                -'av_max':       the density weighted average location of the maxima along one direction
                                -'average':     the average location as weighted by the density profile
                                                along one direction
                                -'point':       center at the point (in fraction of cube size) given by the point kwarg
                                -'pixel':       center at the point (in cube pixels) given by the point kwarg
        :arg str out:       The "out mode". One of the modes accepted by fc._returner. Possible values are
                            'inplace':  save the ndarray data to self.cube return the FractalCube instance;
                            'ndarray':  save the ndarray data to self.cube and return the ndarray;
                            'copy':     return a copy of the FractalCube instance, with the ndarray data saved to the
                                        copy's self.cube.
        :arg bool return_point:    If True, point is also returned.

        :arg: tup point     A three-tuple where the center should be. Only if mode=='point' or mode=='pixel',
                            ignored otherwise.

        :return:            Centered cube. If point==True, result is returned as cube central point is also returned.
        :rtype:             Type depends on 'out'. By default it returns a copy.

        """

        # Get the translation deltas first according to mode
        if mode == 'max':
            max_point = np.nanargmax(fc.cube)
            max_point = np.unravel_index(max_point, dims=fc.shape)
            if ax:
                delta = [0, 0, 0]
                delta[ax] = np.float(max_point[ax]) / fc.shape[ax]
            else:
                delta = list(np.array(max_point, dtype=np.float) / np.array(fc.shape))

        elif mode == 'mean_max':
            if ax:
                max_points = np.nanargmax(fc.cube, axis=ax)
                max_point = np.round(np.nanmean(max_points))
                delta = [0, 0, 0]
                delta[ax] = np.float(max_point[ax]) / fc.shape[ax]
            else:
                max_point = [np.round(np.nanmean(np.nanargmax(fc.cube, axis=0))),
                             np.round(np.nanmean(np.nanargmax(fc.cube, axis=1))),
                             np.round(np.nanmean(np.nanargmax(fc.cube, axis=2)))]
                delta = np.array(max_point, dtype=np.float) / np.array(fc.shape)

        elif mode == 'av_max':
            if ax:
                max_points = np.nanargmax(fc.cube, axis=ax)
                max_vals = np.nanmax(fc.cube, axis=ax)
                max_point = np.round(np.average(max_points, weights=max_vals))
                delta = [0, 0, 0]
                delta[ax] = np.float(max_point) / fc.shape[ax]
            else:
                max_point = []
                for itr in range(3):
                    max_points = np.nanargmax(fc.cube, axis=itr)
                    max_vals = np.nanmax(fc.cube, axis=itr)
                    max_point_itr = np.round(np.average(max_points, weights=max_vals))
                    max_point.append(max_point_itr)
                delta = np.array(max_point, dtype=np.float) / np.array(fc.shape)

        elif mode == 'average':
            if ax:
                rolled_fcshape = np.roll(fc.shape, 3 - ax)
                indexes = np.ones(rolled_fcshape) * np.arange(rolled_fcshape[0])[:, None, None]
                # TODO: check if axis=2 is correct here
                max_points = np.average(indexes, weights=fc.cube, axis=2)
                max_vals = np.nanmax(fc.cube, axis=ax)
                max_point = np.round(np.average(max_points, weights=max_vals))
                delta = [0, 0, 0]
                delta[ax] = np.float(max_point[ax]) / fc.shape[ax]
            else:
                max_point = []
                for itr in range(3):
                    rolled_fcshape = np.roll(fc.shape, 3 - itr)
                    indexes = np.ones(rolled_fcshape) * np.arange(rolled_fcshape[0])[:, None, None]
                    max_points = np.average(indexes, weights=fc.cube, axis=2)
                    max_vals = np.nanmax(fc.cube, axis=itr)
                    max_point_itr = np.round(np.average(max_points, weights=max_vals))
                    max_point.append(max_point_itr)
                delta = np.array(max_point, dtype=np.float) / np.array(fc.shape)

        elif mode == 'point':
            delta = point

        elif mode == 'pixel':
            delta = np.array(point, type=np.float) / np.array(fc.shape)

        if return_point: rp = delta

        delta = - np.array(delta, dtype=np.float) + np.array([0.5, 0.5, 0.5])

        # Perform the translations
        if ax:
            result = self.translate(fc, ax, delta[ax], out='ndarray')
        else:
            fc_copy = self.translate(fc, 0, delta[0], out='copy')
            fc_copy = self.translate(fc_copy, 1, delta[1], out='copy')
            result = self.translate(fc_copy, 2, delta[2], out='ndarray')

        if return_point:
            return fc._returner(result, out), rp
        else:
            return fc._returner(result, out)


class FCExtractor():
    """
    This class contains methods to perform cube extractions of a fractal cube.
    """

    # TODO program trim option

    def __init__(self, bg="small"):

        # Capture the fractal cube object
        self.__dict__.update(locals());
        del self.__dict__['self']

        # ways to order extracted features
        self.orders = ['count', 'sum', 'mean', 'var', 'label']

    def extract_feature(self, fc, low=1.0, order='count', rank=0,
                        bgv=1.e-12, trim=False, out='copy'):
        """ 
        Uses np.ndimage.measurements.label to "label" features in fractal cube.

        :arg obj fc:         :class:`pyFC.FractalCube` object.
        :arg flt low:        Lower threshold limit for extraction
        :arg str order:      One of 'size', 'sum', 'var', 'extent', 'label'
                             Sets which order the features are given in and chosen
                             with *rank*:

                             -'count':     size by number of pixels
                             -'sum':       the sum of the values in the features
                             -'mean':      size by mean of the values in the features
                             -'var':       the varianace in the features
                             -'label':     the default labelling

        :arg int rank:       Which feature will be selected. The ordering is given by *mode*
        :arg flt bgv:        value for the pixels that are not the feature of the final
                             output. Default is a small number, 1.e-12.
        :arg bool trim:      False returns result with same dimensions as input fractal
                             cube. True returns the trimmed result from
                             ndimage.measurements.find_objects. (Not yet programmed.)
        :arg str out:        One of the modes accepted by fc._returner.

        :return:             The extracted feature
        :rtype:              Type depends on 'out'. By default it returns a copy
                             of the :class:`pyFC.FractalCube` object.

        """

        import scipy.ndimage.measurements as snm

        # Threshold array to background values of zero 
        #  and identify (label) all features
        lthr = self.lthreshold(fc, low, bgv=0, out='ndarray')
        labelled, nlabels = snm.label(lthr)

        # Perform operation on all features with label comprehension
        # Also, patch numpy counting function name

        assert (order in self.orders)

        if order == 'label':
            index = rank

        else:
            if order == 'count': order += '_nonzero'
            lbl_range = np.arange(1, nlabels + 1)
            sizes = snm.labeled_comprehension(labelled, labelled, lbl_range,
                                              getattr(np, order), int, 0)

            # Get relevant index and create resulting cube
            index = np.argsort(sizes)[::-1][rank] + 1

        # Create resulting ndarray
        result = np.where(np.equal(labelled, index), lthr, bgv)

        return fc._returner(result, out)

    def lthreshold(self, fc, low=1.0, bgv=0, out='copy'):
        """ 
        Apply lower value threshold to fractal cube 

        :arg obj fc:        :class:`pyFC.FractalCube` object.
        :arg flt low:       Lower threshold limit for extraction
        :arg flt bgv:       A value for the background (the points for which 
                            the values were < low) 
        :arg str out:       One of the modes accepted by fc._returner.

        :return:            The thresholded cube.
        :rtype:             Type depends on 'out'. By default it returns a copy
                            of the :class:`pyFC.FractalCube` object.

        """

        result = np.where(np.less(fc.cube, low), bgv, fc.cube)
        return fc._returner(result, out)

    def touches_boundary(self, fc, bgv=1.e-12):
        """
        Returns true if at least one cell along the boundary of a cube
        has value >=bgv
        """
        boundary = np.dstack((fc.cube[:, :, 0],
                              fc.cube[:, :, -1],
                              fc.cube[0, :, :],
                              fc.cube[-1, :, :],
                              fc.cube[:, 0, :],
                              fc.cube[:, -1, :]))

        if np.any(np.greater(boundary, bgv)):
            return True
        else:
            return False


class FCRayTracer():
    """ 
    This class performs raytrace operations
    """

    def __init__(self):
        pass

    def pp_raytrace(self, fc, absorb_coeff=1., emiss_coeff=1., ax=0):
        """
        Plane-parallel raytracer. This routine merely solves the radiative transfer
        integral from the back to the front of the box. We assume that:

        0. That we are integrating along axis=0.
        1. The box size is 1
        2. cell width are uniform and are a/ni
        3. there is no background source
        4. the emissivity is proportional to the cube variable
        5. the absorption coefficient is proportional to the cube variable

        This raytracer is merely implemented for visualization purposes, which
        is the reason we have assumptions 4 and 5.

        :arg obj fc:                :class:`pyFC.FractalCube` object.
        :arg flt absorb_coeff:      Absorption coefficient   
        :arg flt emiss_coeff:       Emission coefficient   
        :arg int ax:                Direction of plane normal. {0|1|2}, 
                                    default is 0 ("y-z plane")

        :return:                    Raytraced image
        :rtype:                     2D ndarray

        """

        # Emission and absorption coefficients
        c_abs = absorb_coeff * fc.cube
        c_emi = emiss_coeff * fc.cube

        # Transmittance from one side of face to each cell
        # Note: c_abs*delta gives optical depths
        # Also make sure transmittance does not include current layer, so
        # the array needs to be shifted by one layer toward the far boundary
        # The near boundary layer should be filled with a plane of 1s.
        total_length = 1.
        if ax == 0: delta = total_length / fc.ni
        if ax == 1: delta = total_length / fc.nj
        if ax == 2: delta = total_length / fc.nk

        transmittance = np.cumprod(np.exp(-c_abs * delta), axis=ax)
        if ax == 0: transmittance = np.append(np.ones((1, fc.nj, fc.nk)),
                                              transmittance[:-1, :, :], axis=ax)
        if ax == 1: transmittance = np.append(np.ones((fc.ni, 1, fc.nk)),
                                              transmittance[:, :-1, :], axis=ax)
        if ax == 2: transmittance = np.append(np.ones((fc.ni, fc.nj, 1)),
                                              transmittance[:, :, :-1], axis=ax)

        # The resulting integrated result as a 2D array
        return np.sum(c_emi * (1. - np.exp(-c_abs * delta)) * transmittance, axis=ax)


class FCDataEditor():
    def __init__(self):
        pass

    def mult(self, fc, factor, out='copy'):
        """
        Multiply the cube with a value.

        :arg obj fc:        :class:`pyFC.FractalCube` object.
        :arg flt factor:    Factor to multiply cube with
        :arg str out:       One of the modes accepted by fc._returner.

        :return:            Multiplied cube
        :rtype:             Type depends on 'out'. By default it returns a copy
                            of the :class:`pyFC.FractalCube` object.
        """
        result = fc.cube * factor
        return fc._returner(result, out)

    def pow(self, fc, power, out='copy'):
        """
        Exponentiate the cube with a value.

        :arg obj fc:        :class:`pyFC.FractalCube` object.
        :arg flt power:     Power to exponentiate cube with
        :arg str out:       One of the modes accepted by fc._returner.

        :return:            Exponentiated cube
        :rtype:             Type depends on 'out'. By default it returns a copy
                            of the :class:`pyFC.FractalCube` object.

        """

        result = (fc.cube) ** power
        return fc._returner(result, out)


class FCStats():
    """
    Simple class that just collects statistical functions.
    useful for random fractal cube statistics.

    The benefit is that the returner function has the option 
    returning the modified fractal cube with the value of the 
    statistical parameter saved as an attribute, out='inplace', 
    a copy of the fractal cube, out='copy', or just the value,
    out='value'

    .. note:: Untested.
    """

    def __init__(self):

        self.__dict__.update(locals())
        del self.__dict__['self']

    def mean(self, fc, out='value'):
        """
        Calculate mean of fractal cube PDF

        :arg obj fc:    Fractal cube object
        :arg str out:   Output mode. Usually 'value'.

        :return:        Mean of cube PDF
        :rtype:         Double

        """
        mean = np.mean(fc.cube)
        return self._stats_returner(fc, out, mean=mean)

    def std(self, fc, out='value'):
        """
        Calculate standard deviation of fractal cube PDF

        :arg obj fc:    Fractal cube object
        :arg str out:   Output mode. Usually 'value'.

        :return:        Standard deviation of cube PDF
        :rtype:         Double

        """
        std = np.std(fc.cube)
        return self._stats_returner(fc, out, std=std)

    def var(self, fc, out='value'):
        """
        Calculate variance of fractal cube PDF

        :arg obj fc:    Fractal cube object
        :arg str out:   Output mode. Usually 'value'.

        :return:        Variance of cube PDF
        :rtype:         Double

        """
        var = np.var(fc.cube)
        return self._stats_returner(fc, out, var=var)

    def median(self, fc, out='value'):
        """
        Calculate median of fractal cube PDF

        :arg obj fc:    Fractal cube object
        :arg str out:   Output mode. Usually 'value'.

        :return:        Median of cube PDF
        :rtype:         Double

        """
        median = np.median(fc.cube)
        return self._stats_returner(fc, out, median=median)

    def rms(self, fc, out='value'):
        """
        Calculate root mean square of fractal cube PDF

        :arg obj fc:    Fractal cube object
        :arg str out:   Output mode. Usually 'value'.

        :return:        RMS of cube PDF
        :rtype:         Double

        """
        rms = np.sqrt(np.mean(fc.cube ** 2))
        return self._stats_returner(fc, out, rms=rms)

    def skew(self, fc, out='value'):
        """
        Calculate skew of fractal cube PDF

        :arg obj fc:    Fractal cube object
        :arg str out:   Output mode. Usually 'value'.

        :return:        Skew of cube PDF
        :rtype:         Double

        """

        # TODO program without scipy
        skew = sps.skew(fc.cube)

        return self._stats_returner(fc, out, skew=skew)

    def kurt(self, fc, out='value'):
        """
        Calculate kurtosis of fractal cube PDF

        :arg obj fc:    Fractal cube object
        :arg str out:   Output mode. Usually 'value'.

        :return:        Kurtosis of cube PDF
        :rtype:         Double

        """

        # TODO program without scipy
        kurt = sps.kurtosis(fc.cube)

        return self._stats_returner(fc, out, kurt=kurt)

    def flatness(self, fc, out='value'):
        """
        Calculate flatness of fractal cube PDF.
        See appendix in Sutherland & Bicknell (2007).

        :arg obj fc:    Fractal cube object
        :arg str out:   Output mode. Usually 'value'.

        :return:        Flatness of cube PDF
        :rtype:         Double

        """

        # TODO implement this
        flat = 0

        return self._stats_returner(fc, out, flat=flat)

    def filling_factor_theoretical(self, fc, nhot=1., nwarm=1000., thot=1.e7, tcrit=3.e4, out='value'):
        """
        Calculate filling factor of fractal cube given a density ratio and a critical density ratio

        :arg obj fc:    Fractal cube object
        :arg str out:   Output mode. Usually 'value'.
        :arg flt nhot:  Hot phase number density
        :arg flt nwarm: Warm phase number density
        :arg flt thot:  Hot phase temperature
        :arg flt tcrit: Critical temperature (usually 3.e4 K)

        :return:        Filling factor of cube PDF
        :rtype:         Double

        """

        ncrit = nhot * thot / tcrit
        s = fc.sigma * nwarm
        u = fc.mu * nwarm
        a = s * s / (u * u) + 1
        ffv = 1. - 0.5 * (1. + spf.erf(np.log(np.sqrt(a) * ncrit / u) / np.sqrt(2. * np.log(a))))
        return self._stats_returner(fc, out, ffv=ffv)

    def _stats_returner(self, fc, out, **kwargs):
        """
        Different output depending on ``out``.

        :arg str out:   Output mode. {'inplace'|'copy'|'value'|'values'}

                        - 'inplace': Sets the value of the stat attribute and returns cube.
                        - 'copy': Copy of fractal cube with calculated attribute value set.
                        - 'value': Return only the value
                        - 'values': No clue what this was supposed to be.

        :returns:       Result
        :rtype:         Depends on out. (See above.)

        :raise:         ValueError if unknown out.
        """

        if out == 'inplace':
            for k, v in kwargs: setattr(fc, k, v)
            return fc

        elif out == 'copy':
            copy = fc.copy()
            for k, v in kwargs: setattr(copy, k, v)
            return copy

        elif out == 'value':
            return kwargs.values()[0]

        elif out == 'values':
            return kwargs.values()

        else:
            print('Invalid output mode')
            raise ValueError

            # TODO: Include filling factor calculation routines (given a density ratio and a critical density ratio)


# TODO: make everything more "functional". Classes should not carry around data, they should just process it.


class FractalCube:
    """
    Base class for fractal cubes. This class contains basic reading, writing functions, copy function, and functions to
    calculate power spectra. The statistical parameters of the fractal cube are contained as members.
    """

    def __init__(self, ni=64, nj=64, nk=64, kmin=1, kmax=None,
                 mean=1, sigma=np.sqrt(5.), beta=-5. / 3.):

        """
        :arg int i,j,k:     number of points in :math:`i`, :math:`j`, :math:`k` directions.
        :arg int kmin:      :math:`k_\mathrm{min}` of cube.
        :arg int kmax:      :math:`k_\mathrm{max}` of cube
        :arg flt mean:      the mean of the single point distribution.
        :arg flt sigma:     the variance of the single point distribution.
        :arg flt beta:      the power of :math:`D(k)`, the power spectrum.

        :var cube:          The fractal cube data itself.
        :var ndim:          dimensionality of cube.
        :var shape:         [ni, nj, nk]

        """

        # Nyquist limit
        nq_limit = np.floor(max([ni, nj, nk]) / 2.)

        # Set Nyquist limit if kmax == None
        # This ensures kmax !> Nyquist limit in all directions since
        # both the Nyquist limit and kmax scale proportionally
        if kmax == None or kmax > nq_limit:
            kmax = nq_limit

        # Assign input values as members.
        self.ni = ni
        self.nj = nj
        self.nk = nk
        self.kmin = kmin
        self.kmax = kmax
        self.mean = mean
        self.sigma = sigma
        self.beta = beta
        self.nq_limit = nq_limit

        # Save some useful values
        self.shape = [ni, nj, nk]
        self.ndim = np.sum(np.greater(self.shape, 1))

        # Dictionary of manipulators (just for bookkeeping)
        # {name of object: [class, functions...]}
        self.manipulators = OrderedDict((
            ['fc_slicer', [FCSlicer, 'slice', 'tri_slice']],
            ['fc_affine', [FCAffine, 'translate', 'permute', 'mirror']],
            ['fc_extractor', [FCExtractor, 'extract_feature', 'lthreshold']],
            ['fc_raytracer', [FCRayTracer, 'pp_raytrace']]
        ))
        self._get_all_manipulators()

        # Obtain box ratios
        norder = np.argsort(self.shape)
        inorder = np.argsort(norder)
        shp_ord = np.array(self.shape)[norder]
        self.cube_ratios = (1.0 * shp_ord / shp_ord[2])[inorder]

        # Obtain maximum values of kmin and kmax for box directions,
        # which are smaller than the largest box dimension.
        # This is not used in cloud generation procedure.
        # Instead, the sampling vector is scaled appropriately.
        # However, kmin !< 1 in any direction after the scaling,
        # and we check this here.
        self.cube_kmins = ckm = kmin * self.cube_ratios
        if np.any(np.logical_and(np.less(ckm, 1.), np.greater(self.shape, 1))):
            print('Error: kmin for at least one dimension is less than 1.')
            print('Increase the number of cells in that dimension, or increase kmin.')
            raise (ValueError)

        # Store kmaxs and ny_limits per direction
        # kmax !> Nyquist limit in any direction is already satisfied
        self.cube_kmaxs = kmax * self.cube_ratios
        self.cube_nq_limits = nq_limit * self.cube_ratios

    def copy(self):
        """
        Make a deep copy of this object

        :return:    A deep copy of this object
        :rtype:     :class:`pyFC.FractalCube`
        """

        return copy.deepcopy(self)

    def read_cube(self, fname='data.dbl', prec='double', out='inplace'):
        """
        Read cube from file. The data of the cube returned and also
        saved to the :attr:`pyFC.LogNormalFractalCube.cube` attribute

        :arg str fname:     Data filename
        :arg str prec:      Floating point precision of output {'double'|'single'}

        :return:            Data of cube read in
        :rtype:             ndarray

        :raise:             ValueError if prec is invalid
        """
        if prec == 'double':
            dtype = np.dtype('<f8')
        elif prec == 'single':
            dtype = np.dtype('<f4')
        else:
            ValueError('Unknown prec ' + prec)

        cube = np.fromfile(fname, dtype=dtype, count=-1)
        cube = cube.reshape(*self.shape)

        return self._returner(cube, out)

    def write_cube(self, fc=None, fname='data.dbl', app=True, prec='double'):
        """
        Writes out a fractal cube data file in little endian, 
        double precision. Care is taken not to overwrite existing files.

        :arg str fname:     Data filename
        :arg bool app:      automatically append kmin, ni, nj, and nk values to filename.
                            If app == True, append kmin, ni, nj, nk values to filename.
                            If suffixed files '<base>-[0-9][0-9]<ext>' exist, appended
                            by a suffix one larger than the highest number found.
        :arg str prec:      Floating point precision of output {'double'|'single'}

        :return:            1 for successful write
        :rtype:             int

        :raise:             ValueError if prec is invalid
        """

        if fc is None:
            cube = self.cube
        else:
            cube = fc.cube

        if prec == 'double':
            out = cube.astype('<f8')
        elif prec == 'single':
            out = cube.astype('<f4')
        else:
            ValueError('Unknown prec ' + prec)

        if app:
            ext = osp.splitext(fname)[1]
            base = osp.splitext(fname)[0]

            fname = (f'{base}_k{self.kmin:0>2d}_n{self.ni:0>4d}' +
                     (f'x{self.nj:0>4d}' if self.ndim > 1 else '') +
                     (f'x{self.nk:0>4d}' if self.ndim > 2 else '') +
                     f'_s{self.sigma:0>4d}' + ext)

        fname = pt.unique_fname(fname, '-', '[0-9][0-9]')

        out.T.tofile(fname)
        return 1

[docs] def _returner(self, result, out): """ Given a result, return the correct thing given an "out mode" :arg ndarray result: The FractalCube.cube result that should be returned, copied, etc. :arg str out: The "out mode". Possible values are 'inplace': save the ndarray data to self.cube return the FractalCube instance; 'ndarray': save the ndarray data to self.cube and return the ndarray; 'copy': return a copy of the FractalCube instance, with the ndarray data saved to the copy's self.cube. :return: Varies according to the "out mode", ``out`` :rtype: Varies according to the "out mode", ``out`` inplace: default, change in place """ if out == 'inplace': self.cube = result return self elif out == 'ndarray': # self.cube = result return result elif out == 'copy': copy = self.copy() copy.cube = result return copy else: print('Invalid output mode. Choose one of "inplace", "ndarray", or "copy". ') raise (ValueError)
def _get_all_manipulators(self): """ Create methods in self from self.manipulators. The manipulator objects are also created and stored in self. """ # Create method attributes. for o, m in self.manipulators.items(): for f in m[1:]: setattr(self, f, self._m_curried(getattr(m[0](), f))) def _m_curried(self, m_func): """ Currying function for creating method attributes whose first argument is the fractal cube, self, itself, rather an external cube instance. """ def m_curry(*args, **kwargs): return m_func(self, *args, **kwargs) return m_curry def power_spec(self, cube=None): """ Power spectrum of a fractal cube data :arg ndarray cube: FractalCube.cube data array :return: Power spectrum of this fractal cube :rtype: ndarray """ if cube is None: cube = self.cube return self._power_spec(fft.fftn(cube)) def iso_power_spec(self, kmag=None, cube=None, stdev=False): """ Isotropic power spectrum of this fractal cube. :arg ndarray kmag: Array of wave-vector magnitude values (same size and dimensions as cube) :arg ndarray cube: Array with cube data :arg bool stdev: Calculate standard deviation :return: Isotropic power spectrum of this fractal cube. :math:`(D(|k|), |k|)` :rtype: tup """ if kmag is None: kmag = self.kmag_sampling() if cube is None: cube = self.cube return self._iso_power_spec(kmag, self._power_spec(fft.fftn(cube)), stdev) def kmag_sampling(self, ni=None, nj=None, nk=None): """ Sampling space, :math:`|k|_{i,j,k}`. See definition of FFT, :math:`k` and frequency at http://docs.scipy.org/doc/numpy/reference/routines.fft.html :arg int ni, nj, nk: Number of cells (sample points) in i, j, k directions, respectively :return: 3-D array with sampling vector magnitudes :rtype: ndarray """ if ni is None: ni = self.ni if nj is None: nj = self.nj if nk is None: nk = self.nk sampli, samplj, samplk = fft.fftfreq(ni), fft.fftfreq(nj), fft.fftfreq(nk) k1di = sampli * ni / self.cube_ratios[0] k1dj = samplj * nj / self.cube_ratios[1] k1dk = samplk * nk / self.cube_ratios[2] ksqri, ksqrj, ksqrk = k1di * k1di, k1dj * k1dj, k1dk * k1dk kmag = np.sqrt(np.add.outer(np.add.outer(ksqri, ksqrj), ksqrk)) return kmag def norm_spec(self, spectrum): """ Normalize a spectrum so that apodizing a random Gaussian cube with the root of spectrum does not change the mean and variance of the Gaussian :arg ndarray spectrum: A 1-, 2-, or 3-dimensional spectrum :return: Normalized spectrum, same shape as input spectrum :rtype: ndarray """ shape, N = spectrum.shape, spectrum.size spectrum = np.ravel(spectrum) csqr = (N - 1.) / np.sum(np.real(spectrum[1:])) spectrum *= csqr spectrum[0] = 1 return spectrum.reshape(shape) def _iso_power_spec(self, kr, pspec_r, raveled=False, stdev=False, digitize=False): """ kr k array pspec_r power spectrum array stdev output stdvs as well digitize output psd as well (even if stdev is false) kr and pspec must be an array of same shape We ravel the arrays because the digitize function requires 1d array If (already) raveled (at input), then assume that kr is 1d and that the first element represents k0 + 1. If not, ravel and remove first elements of both kr and pspec. """ if not raveled: kr = np.ravel(kr) pspec_r = np.ravel(pspec_r) bins = np.append(np.unique(kr), kr.max() + 1) psc, bins = np.histogram(kr, bins) psw, bins = np.histogram(kr, bins, weights=pspec_r) means = psw / psc binc = 0.5 * (bins[:-1] + bins[1:]) if stdev or digitize: psd = np.digitize(kr, bins[1:-1]) if stdev: pss, bins = np.histogram(kr, bins, weights=(pspec_r - means[psd]) ** 2) stdvs = np.sqrt(pss / psc) if digitize: return means, binc, stdvs, psd.reshape(self.shape) else: return means, binc, stdvs else: if digitize: return means, binc, psd.reshape(self.shape) else: return means, binc def _power_spec(self, F, k=None): """ Power spectrum. Without the factor of 4*pi. """ return np.real(np.conjugate(F) * F) def print_stats(self): """ Prints stats of the Fractal cube, including (cubetype, ni, nj, nk, kmin, mean, sigma, beta) :return: 1 at completion of routine :rtype: int """ # TODO Finish programming this param = [self.cubetype, self.ni, self.nj, self.nk, self.kmin, self.mean, self.sigma, self.beta] param_str = locals().keys()[1:] for i in zip(param_str, param): print(param_str + str(param)) return 1 def inner_range_slice(self, cf, min, max): """ Return a numpy slice object that only selects some region interior of min and max :param cf: values to be compared with min and max :param min: minimum value :param max: maximum value :return: slice object for interior regions """ sf = np.logical_and(np.greater_equal(cf, min), np.less_equal(cf, max)) sf = np.s_[sf] return sf
[docs]class GaussianFractalCube(FractalCube): """ Class for lognormal fractal cubes, as generated by taking a random gaussian field and apodizing its spectrum in fourier space with a power law. """ # TODO refactor mean to be mu def __init__(self, ni=64, nj=64, nk=64, kmin=1, kmax=None, mean=0, sigma=200., beta=-1.6666666666666): """ :arg int ni,nj,nk: number of points in :math:`i`, :math:`j`, :math:`k` directions :arg int kmin: :math:`k_\mathrm{min}` of fractal cube :arg flt mean: the mean of the lognormal PDF :arg flt sigma: the standard deviation of the gaussian PDF :arg flt beta: the power-law index of :math:`D(k)`, the power spectrum. The above parameters are all saved as attributes of an instance of this class. """ #: (*str*) -- Cube type {'lognormal', 'gaussian'}. self.cubetype = 'gaussian' #: (*ndarray*) -- Gaussian fractal cube data as ndarray of shape :attr:`pyFC.GaussianFractalCube.shape` self.cube = None #: (*int*) -- Number of dimensions of the fractal cube data array :attr:`pyFC.GaussianFractalCube.cube` self.ndim = None #: (*tup*) -- Shape of the fractal cube data array :attr:`pyFC.GaussianFractalCube.cube` self.shape = None FractalCube.__init__(self, ni=ni, nj=nj, nk=nk, kmin=kmin, kmax=kmax, mean=mean, sigma=sigma, beta=beta)
[docs] def func_target_spec(self, k, kmin=None, kmax=None, beta=None): """ The target spectrum function :arg ndarray k: A domain in :math:`k`-space to evaluate the target spectrum over. :arg flt kmin: A lower cutoff wavenumber, :math:`k_\mathrm{min}`. :arg flt beta: The power law index of the target spectrum, :math:`\beta`. :return: Target spectrum. Same shape as k. :rtype: ndarray """ if kmin is None: kmin = self.kmin if kmax is None: kmax = self.kmax if beta is None: beta = self.beta return np.where(np.logical_and(np.greater_equal(np.abs(k), kmin), np.less_equal(np.abs(k), kmax)), np.abs(k) ** (beta - 2.), 0)
# return np.where(np.greater_equal(np.abs(k), kmin), np.abs(k) ** (beta - 2.), 0)
[docs] def gen_cube(self, history=False): """ Generate a Gaussian fractal cube with the statistics saved as attributes in an instance of this object. :arg bool history: If history is True, this routine returns the initial cube too. :return: Gaussian fractal cube :rtype: :class:`pyFC.GaussianFractalCube` """ # Gaussian random field, and the corresponding lognormal random field grf = np.random.normal(self.mean, self.sigma, self.shape) # Get the ndim kmag array (isotropic k) kmag = self.kmag_sampling() # The target spectrum target_spec = self.func_target_spec(kmag) target_spec = self.norm_spec(target_spec) # The apodization values apod = np.sqrt(target_spec) # N-dimensional (3-dimensional) FFT log-normal cube Fg = fft.fftn(grf) # Apodize with power law Fga = Fg * apod # Fourier transform back (the imaginary part is negligible) cube = np.real(fft.ifftn(Fga)) # Return cubes. In "history mode", the cubes in the # beginning is also returned. The history mode is only # used through the routine history_cubes if history: return self._returner(grf, 'copy'), self._returner(cube, 'inplace') else: return self._returner(cube, 'inplace')
[docs] def lnfc(self): """ Corresponding Lognormal field. :return: The corresponding Lognormal field. :rtype: :class:`pyFC.LogNormalFractalCube` object """ ln = mt.LogNormalPDF(self.mean, self.sigma, gstat=True) lnfc = LogNormalFractalCube(self.ni, self.nj, self.nk, ln.mu, ln.sigma, self.beta) lnfc.cube = np.exp(self.cube) return lnfc
[docs]class LogNormalFractalCube(FractalCube): """ Class for lognormal fractal cubes, as generated by the method of `Lewis & Austin (2002) <https://ams.confex.com/ams/11AR11CP/techprogram/paper_42772.htm>`_. This class contains basic reading, writing, generating functions, copy function, and functions to calculate power spectra. Functions to calculate the power spectra exist because they are needed in the actual construction routine. The statistical parameters of the fractal cube are contained as attributes. """ # TODO refactor mean to be mu def __init__(self, ni=64, nj=64, nk=64, kmin=1, kmax=None, mean=1, sigma=np.sqrt(5.), beta=-1.66666666666): """ :arg int ni,nj,nk: number of points in :math:`i`, :math:`j`, :math:`k` directions :arg int kmin: :math:`k_\mathrm{min}` of fractal cube :arg int kmax: :math:`k_\mathrm{max}` of fractal cube. This is set to the Nyquist limit if None. :arg flt mean: the mean of the lognormal PDF :arg flt sigma: the standard deviation of the lognormal PDF :arg flt beta: the power-law index of :math:`D(k)`, the power spectrum. The above parameters are all saved as attributes of an instance of this class. """ # Documented members and their defaults. #: (*str*) -- Cube type {'lognormal', 'gaussian'}. self.cubetype = 'lognormal' #: (*ndarray*) -- The lognormal fractal cube data as an arrayof shape :attr:`pyFC.LogNormalFractalCube.shape` self.cube = None #: (*int*) -- Number of dimensions of the fractal cube data array :attr:`pyFC.LogNormalFractalCube.cube` self.ndim = None #: (*tup*) -- Shape of the fractal cube data array :attr:`pyFC.LogNormalFractalCube.cube` self.shape = None FractalCube.__init__(self, ni=ni, nj=nj, nk=nk, kmin=kmin, kmax=kmax, mean=mean, sigma=sigma, beta=beta)
[docs] def func_target_spec(self, k, kmin=None, kmax=None, beta=None): """ The target spectrum function :arg ndarray k: A domain in :math:`k`-space to evaluate the target spectrum over. :arg flt kmin: A lower cutoff wavenumber, :math:`k_\mathrm{min}`. :arg flt beta: The power law index of the target spectrum, :math:`\beta`. :return: Target spectrum. Same shape as k. :rtype: ndarray """ if kmin is None: kmin = self.kmin if kmax is None: kmax = self.kmax if beta is None: beta = self.beta # Doesn't work yet # return np.where(np.logical_and(np.greater_equal(np.abs(k), kmin), # np.less_equal(np.abs(k), kmax)), # np.abs(k) ** (beta - 2.), 0) return np.where(np.greater_equal(np.abs(k), kmin), np.abs(k) ** (beta - 2.), 0)
[docs] def gen_cube(self, verbose=True): """ Generate a lognormal fractal cube by the method of `Lewis & Austin (2002) <https://ams.confex.com/ams/11AR11CP/techprogram/paper_42772.htm>`_ with the statistical properties saved in the instance of the :class:`pyFC.FractalCube` object. :arg bool verbose: Verbosity switch. :return: The generated lognormal fractal cube :rtype: :class:`pyFC.LogNormalFractalCube` object .. note:: The actual code for generating fractal cubes is in :meth:`pyFC.LogNormalFractalCube.yield_cubes`. """ for fc in self.yield_cubes(history=False, verbose=verbose): return fc
[docs] def yield_cubes(self, history=True, verbose=True): """ Generate a lognormal fractal cube by the method of `Lewis & Austin (2002) <https://ams.confex.com/ams/11AR11CP/techprogram/paper_42772.htm>`_ with the statistical properties saved in the instance of the FractalCube object. :arg bool verbose: Verbosity switch. :arg bool history: History switch. If true, generate a cube for each iteration. This parameter should always be true, unless this function is called from :meth:`pyFC.LogNormalFractalCube.gen_cube`. :return: Yields lognormal fractal cubes :rtype: :class:`pyFC.LogNormalFractalCube` object If history is True, this routine yields the cubes for each iteration. This option is actually not to be used. Use gen_cube instead. """ # _____________________________________ # Parameters (These can be changed) # Iteration self.iter_max = 10 self.iter_tol = 0.01 # Correction factor self.eta = 0.5 # _____________________________________ # Main bit # Theoretical lognormal object ln = mt.LogNormalPDF(self.mean, self.sigma) # Convert to gaussian stats mean_g, sigma_g = ln.mu_g, ln.sigma_g # Gaussian random field, and the corresponding lognormal random field cube = np.random.normal(mean_g, sigma_g, self.shape) # Yield cube, if history is on if history: yield self._returner(np.exp(cube), 'copy') # Get the ndim kmag array (isotropic k) kmag = self.kmag_sampling(self.ni, self.nj, self.nk) # Isotropic k (kmag is also used as a dummy second argument) dummy, k_iso = self._iso_power_spec(kmag, kmag) # Some helpful arrays for log k lkmin, lkmax = np.log10(self.kmin), np.log10(self.kmax) # Some helpful indices that determine where fits and # corrections are applied an. These could have been defined in # linear space, but this doesn't change the indexing. The important # thing is that there is one for k_iso and one for kmag. sf refers to # indices for fitting region, so, for the other region. sf_lk_iso = self.inner_range_slice(mt.zero_log10(k_iso), lkmin, lkmax) sf_lk_iso[0] = False sf_lkmag = self.inner_range_slice(mt.zero_log10(kmag), lkmin, lkmax) sf_lkmag = np.ravel(sf_lkmag) sf_lkmag[0] = False sf_lkmag = sf_lkmag.reshape(self.shape) # The target spectrum target_spec = self.func_target_spec(kmag) target_spec = self.norm_spec(target_spec) target_spec_iso = self.func_target_spec(k_iso) target_spec_iso = mt.zero_log10(target_spec_iso[sf_lk_iso]) # The apodization values target_spec = np.sqrt(target_spec) # N-dimensional (3-dimensional) FFT lognormal cube cube = fft.fftn(cube) # Apodize with power law cube_a = cube * target_spec # Previously cube = Fg, and Fga = cube_a # Loop begins here convergence, iiter = 1, 1 while convergence > self.iter_tol and iiter <= self.iter_max: # Fourier transform back (the imaginary part is negligible) cube_a = np.real(fft.ifftn(cube_a)) # Create lognormal cube_a = np.exp(cube_a) # Yield cube, if history if history: yield self._returner(cube_a, 'copy') # Power spectrum of lognormal is not desired power-law cube_a = fft.fftn(cube_a) cube_a = self._power_spec(cube_a) cube_a, k_iso = self._iso_power_spec(kmag, cube_a) # Fits to the isotropic spectra cube_a = mt.zero_log10(cube_a[sf_lk_iso]) # Zeroth order fit, to find best height for target spectrum # (kind of normalization). The multiplication by 10**fit0 # retains zeroes in log space. fit0 = np.average(cube_a - target_spec_iso, weights=np.r_[np.diff(k_iso), np.diff(k_iso[-2:])][sf_lk_iso]) # Fit power spec of lognormal with second order polynomial fit2 = np.polyfit(mt.zero_log10(k_iso)[sf_lk_iso], cube_a, 2) # Corrections based on fits. fit0 needs to be multiplied # with the func_target_spec, otherwise 0s are not preserved. p2 = np.polyval(fit2, mt.zero_log10(kmag)) p0 = mt.zero_log10(10 ** fit0 * self.func_target_spec(kmag)) corr = np.where(sf_lkmag, self.eta * (p0 - p2), 0) # Apply correction (in log space) to apodization spectrum corr = target_spec * 10 ** corr # Re-Apodize with normalized power law # From here, it's the same as before the loop. corr = self.norm_spec(corr * corr) target_spec_old = target_spec.copy() target_spec = np.sqrt(corr) cube_a = cube * target_spec # Estimate convergence convergence = np.average(mt.zero_div(abs(self._power_spec(target_spec_old) - self._power_spec(target_spec)), self._power_spec(target_spec))) # Some messages if verbose: print('iteration ' + str(iiter)) print('convergence = ' + str(convergence)) print('') # Get ready for next iteration iiter += 1 # A last conversion # Fourier transform back (the imaginary part is negligible) cube = np.real(fft.ifftn(cube_a)) # Create lognormal cube = np.exp(cube) # Return cube object yield self._returner(cube, 'inplace')
[docs] def gnfc(self): """ The corresponding Gaussian field. :return: The corresponding Gaussian field. :rtype: :class:`pyFC.GaussianFractalCube` object """ ln = mt.LogNormalPDF(self.mean, self.sigma) gnfc = GaussianFractalCube(self.ni, self.nj, self.nk, self.kmin, self.kmax, ln.mu_g, ln.sigma_g, self.beta) gnfc.cube = np.log(self.cube) return gnfc