dragon_scales

#!/usr/bin/env python
#
# Copyright (C) 2014 Smithsonian Astrophysical Observatory
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#

from __future__ import print_function

toolname = "dragon_scales"
__revision__ = "15 January 2019"

import os
import sys
import numpy as np

import ciao_contrib.logger_wrapper as lw
lgr = lw.initialize_logger(toolname)
verb0 = lgr.verbose0
verb1 = lgr.verbose1
verb2 = lgr.verbose2
verb3 = lgr.verbose3
verb5 = lgr.verbose5

from cxcdm import *



class OverlappingShapes():

    invalid_pixel = 12345678
    
    def __init__(self, infile, psffile, gradient ):
        """
        """
        from crates_contrib.masked_image_crate import MaskedIMAGECrate
        self.img = MaskedIMAGECrate(infile,mode="r")

        self.vals = self.img.get_image().values
        self.infile = infile

        try:
            self.pmap = np.zeros_like( self.vals ) + float( psffile )
        except Exception as e:
            self.pmap = MaskedIMAGECrate( psffile, mode="r").get_image().values

        self.nullval = -999
        self.mask = None

        if gradient:
            gx,gy = np.gradient( self.vals)
            self.angle = np.arctan2(gy,gx)  # radians, keep it like that
        else:
            self.angle = np.zeros_like(self.vals)


    def get_psfsize( self, xx, yy):
        """
        """
        import numpy as np
        p = self.pmap[yy,xx]
        return p if np.isfinite(p) else 0

    @staticmethod
    def inside( dx, dy, rad ):
        raise NotImplementedError("Please implement this in the subclass")
        
    def set_mask( self, xx, yy, maskval ):
        """
        """

        rad = int(self.get_psfsize( xx, yy )+0.5)
        
        rad_2 = 2*rad # allow for rotation of shapes
        

        def rotate( _x, _y ):
            aa = -self.angle[yy,xx]
            cos_a = np.cos(aa)
            sin_a = np.sin(aa)        
            ox =  cos_a * _x + sin_a * _y
            oy = -sin_a * _x + cos_a * _y
            return(ox,oy)
        

        for _y in range( yy-rad_2-1, yy+rad_2+1, 1):
            if _y < 0 or _y >= self.mask.shape[0]:
                continue
            dy = _y-yy
            for _x in range( xx-rad_2-1, xx+rad_2+1, 1):
                if _x < 0 or _x >= self.mask.shape[1]:
                    continue
                if self.mask[_y,_x] != 0:
                    continue            
                if not self.img.valid(_x,_y):
                    self.mask[_y, _x] = self.invalid_pixel
                    self.vals[_y, _x] = self.nullval
                    continue

                dx = _x-xx

                rx,ry = rotate(dx,dy)
                if self.inside( rx,ry, rad):
                    self.mask[_y, _x] = maskval
                    self.vals[_y, _x] = self.nullval


    def get_loc_max(self):
        """
        """
        ll = np.where( self.vals == np.max( self.vals ) )
        ly = ll[0][0]
        lx = ll[1][0]
        
        self.vals[ly,lx] = self.nullval
        return (ly,lx)


    def doit(self):
        """
        """
        self.mask = np.zeros_like(self.vals).astype(np.int32)

        maskval = 0
        while 0 == np.min(self.mask):
            maskval = maskval + 1
            verb3(str(maskval))
            ly,lx = self.get_loc_max()
            self.set_mask( lx, ly, maskval )


        for iy in range(self.mask.shape[0]):
            for ix in range(self.mask.shape[1]):
                if self.mask[iy][ix] == self.invalid_pixel:
                    self.mask[iy][ix] = 0

            
    def save( self, outfile ):
        """
        """
        if os.path.exists( outfile ):
            os.remove( outfile)  # clobber check already done

        dm_img = dmImageOpen(self.infile+"[opt type=i4]")
        out_img = dmBlockCreateCopy( dmDatasetCreate(outfile), "GRID", dm_img)
        out_dd = dmImageGetDataDescriptor( out_img )
        dmImageSetData( out_dd, self.mask )
        dmImageClose( out_img )
        dmImageClose( dm_img)


class RegionTiles( OverlappingShapes ):
    """
    Use cxcRegion for shapes
    """
    def make_region(self, radius):
        raise NotImplementedError("Implement in the subclass")

            
    def inside( self, dx, dy, rad ):        
        if not hasattr(self, "region"):
            self.make_region(rad)
                
        return self.region.is_inside( dx,dy)


class HexTiles( RegionTiles):
    """
    
    """
    def make_region(self, radius):
        import region as cxcregion
        s60 = np.sin(np.deg2rad(60.0))
        c60 = np.cos(np.deg2rad(60.0))
        px = np.array( [ -1.0, -c60, c60, 1.0,  c60, -c60, -1.0 ] )*radius
        py = np.array( [ 0.0,   s60, s60, 0.0, -s60, -s60,  0.0 ] )*radius
        self.region = cxcregion.polygon(px,py)


class TriangleTiles( RegionTiles):
    """
    
    """    
    def make_region(self, radius):
        import region as cxcregion
        px = np.array( [ -0.5, 0.5, 0 ] )*radius
        py = np.array( [ -0.5, -0.5, 0.5 ] )*radius
        self.region = cxcregion.polygon(px,py)

    
class InvertedTriangleTiles( RegionTiles):
    """
    
    """    
    def make_region(self, radius):
        import region as cxcregion
        px = np.array( [ -0.5, 0.5, 0 ] )*radius
        py = np.array( [ 0.5, 0.5, -0.5 ] )*radius
        self.region = cxcregion.polygon(px,py)

    
class RoofingTiles( OverlappingShapes ):
    """
    Use square shapes
    """
    @staticmethod
    def inside( dx, dy, rad ):        
        return (max( [abs(dx),abs(dy)] ) <= rad)


class DragonScales( OverlappingShapes ):
    """
    Use circular shapes
    """
    @staticmethod
    def inside( dx, dy, rad ):
        from math import hypot
        return (hypot( dx, dy) <= rad)

    
class Diamonds( OverlappingShapes ):
    """
    Use city-block distance
    """
    @staticmethod
    def inside( dx, dy, rad ):
        return (abs(dx)+abs(dy) <= rad)

   
@lw.handle_ciao_errors( toolname, __revision__) 
def main():
    """
    
    """

    from ciao_contrib.param_soaker import get_params
    from ciao_contrib.runtool import dmmaskbin
    from ciao_contrib.runtool import add_tool_history

    # Load parameters
    pars = get_params(toolname, "rw", sys.argv, 
        verbose={"set":lw.set_verbosity, "cmd":verb1} )

    from ciao_contrib._tools.fileio import outfile_clobber_checks
    outfile_clobber_checks(pars["clobber"], pars["outfile"] )
    outfile_clobber_checks(pars["clobber"], pars["binimg"] )
    gradient = ("yes" == pars["gradient"])

    if "circle" == pars["shape"]:
        dragon = DragonScales( pars["infile"], pars["radius"], gradient )
    elif "box" == pars["shape"]:
        dragon = RoofingTiles(  pars["infile"], pars["radius"], gradient )
    elif "diamond" == pars["shape"]:
        dragon = Diamonds(  pars["infile"], pars["radius"], gradient )
    elif "hexagon" == pars["shape"]:
        dragon = HexTiles(  pars["infile"], pars["radius"], gradient )
    elif "triangle" == pars["shape"]:
        dragon = TriangleTiles(  pars["infile"], pars["radius"], gradient )
    elif "itriangle" == pars["shape"]:
        dragon = InvertedTriangleTiles(  pars["infile"], pars["radius"], gradient )
    else:
        raise ValueError("Unknown shape value {}".format(pars["shape"]))

    dragon.doit()
    dragon.save( pars["outfile"] )
    add_tool_history( pars["outfile"], toolname, pars, toolversion=__revision__)

    if len(pars["binimg"])>0 and "none" != pars["binimg"].lower():
        dmmaskbin( pars["infile"], pars["outfile"]+"[opt type=i4]", pars["binimg"], clobber=True)
        add_tool_history( pars["binimg"], toolname, pars, toolversion=__revision__)

    


if __name__ == "__main__":
    try:
        main()
    except Exception as E:
        print("\n# "+toolname+" ("+__revision__+"): ERROR "+str(E)+"\n", file=sys.stderr)
        sys.exit(1)
    sys.exit(0)