Add 'region.py' to manipulate regions in CIAO format

1 files changed, 251 insertions, 0 deletions

diff --git a/scripts/region.py b/scripts/region.py
new file mode 100644
index 0000000..09e42fa
--- /dev/null
+++ b/scripts/region.py
@@ -0,0 +1,251 @@
+# Copyright (c) 2017 Weitian LI <liweitianux@live.com>
+# MIT license
+
+"""
+Region utilities.
+
+NOTE:
+Only support the commonly used region shapes in CIAO format.
+While more complex shapes and DS9 format are not considered
+at the moment.
+"""
+
+import os
+import re
+
+
+class Point:
+    """
+    Point region: point(xc,yc)
+    """
+    shape = "point"
+
+    def __init__(self, xc=None, yc=None):
+        if isinstance(xc, str):
+            self.parse(regstr=xc)
+        else:
+            self.xc = xc
+            self.yc = yc
+
+    def parse(self, regstr):
+        g = re.split(r"[(),]", regstr)
+        self.xc = float(g[1])
+        self.yc = float(g[2])
+
+    def __str__(self):
+        return "%s(%.4f,%.4f)" % (self.shape, self.xc, self.yc)
+
+
+class Circle:
+    """
+    Circle region: circle(xc,yc,radius)
+    """
+    shape = "circle"
+
+    def __init__(self, xc=None, yc=None, radius=None):
+        if isinstance(xc, str):
+            self.parse(regstr=xc)
+        else:
+            self.xc = xc
+            self.yc = yc
+            self.radius = radius
+
+    def parse(self, regstr):
+        g = re.split(r"[(),]", regstr)
+        self.xc = float(g[1])
+        self.yc = float(g[2])
+        self.radius = float(g[3])
+
+    def __str__(self):
+        return "%s(%.4f,%.4f,%.4f)" % (self.shape, self.xc, self.yc,
+                                       self.radius)
+
+
+class Annulus:
+    """
+    Annulus region:
+        annulus(xc,yc,radius,radius2)
+        + xc, yc : center
+        + radius, radius2 : inner, outer radius
+    """
+    shape = "annulus"
+
+    def __init__(self, xc=None, yc=None, radius=None, radius2=None):
+        if isinstance(xc, str):
+            self.parse(regstr=xc)
+        else:
+            self.xc = xc
+            self.yc = yc
+            self.radius = radius
+            self.radius2 = radius2
+
+    def parse(self, regstr):
+        g = re.split(r"[(),]", regstr)
+        self.xc = float(g[1])
+        self.yc = float(g[2])
+        self.radius = float(g[3])
+        self.radius2 = float(g[4])
+
+    def __str__(self):
+        return "%s(%.4f,%.4f,%.4f,%.4f)" % (self.shape, self.xc, self.yc,
+                                            self.radius, self.radius2)
+
+
+class Pie:
+    """
+    Pie region:
+        pie(xc,yc,radius,radius2,angle1,angle2)
+        + xc, yc : center
+        + radius, radius2 : inner, outer radius
+        + angle1, angle2: start, end angle [0, 360)
+    """
+    shape = "pie"
+
+    def __init__(self, xc=None, yc=None, radius=None, radius2=None,
+                 angle1=0.0, angle2=360.0):
+        if isinstance(xc, str):
+            self.parse(regstr=xc)
+        else:
+            self.xc = xc
+            self.yc = yc
+            self.radius = radius
+            self.radius2 = radius2
+            self.angle1 = angle1
+            self.angle2 = angle2
+
+    def parse(self, regstr):
+        g = re.split(r"[(),]", regstr)
+        self.xc = float(g[1])
+        self.yc = float(g[2])
+        self.radius = float(g[3])
+        self.radius2 = float(g[4])
+        self.angle1 = float(g[5])
+        self.angle2 = float(g[6])
+
+    def __str__(self):
+        return "%s(%.4f,%.4f,%.4f,%.4f,%.4f,%.4f)" % (
+            self.shape, self.xc, self.yc, self.radius, self.radius2,
+            self.angle1, self.angle2)
+
+
+class Ellipse:
+    """
+    Ellipse region:
+        ellipse(xc,yc,radius,radius2,rotation)
+        + xc, yc : center
+        + radius, radius2 : semi-major / semi-minor axis
+        + rotation: rotation angle [0, 360)
+    """
+    shape = "ellipse"
+
+    def __init__(self, xc=None, yc=None, radius=None, radius2=None,
+                 rotation=0.0):
+        if isinstance(xc, str):
+            self.parse(regstr=xc)
+        else:
+            self.xc = xc
+            self.yc = yc
+            self.radius = radius
+            self.radius2 = radius2
+            self.rotation = rotation
+
+    def parse(self, regstr):
+        g = re.split(r"[(),]", regstr)
+        self.xc = float(g[1])
+        self.yc = float(g[2])
+        self.radius = float(g[3])
+        self.radius2 = float(g[4])
+        self.rotation = float(g[5])
+
+    def __str__(self):
+        return "%s(%.4f,%.4f,%.4f,%.4f,%.4f)" % (
+            self.shape, self.xc, self.yc, self.radius,
+            self.radius2, self.rotation)
+
+
+class Box:
+    """
+    Box region: box(xc,yc,width,height,rotation)
+    """
+    shape = "box"
+
+    def __init__(self, xc=None, yc=None, width=None, height=None,
+                 rotation=0.0):
+        if isinstance(xc, str):
+            self.parse(regstr=xc)
+        else:
+            self.xc = xc
+            self.yc = yc
+            self.width = width
+            self.height = height
+            self.rotation = rotation
+
+    def parse(self, regstr):
+        g = re.split(r"[(),]", regstr)
+        self.xc = float(g[1])
+        self.yc = float(g[2])
+        self.width = float(g[3])
+        self.height = float(g[4])
+        self.rotation = float(g[5])
+
+    def __str__(self):
+        return "%s(%.4f,%.4f,%.4f,%.4f,%.4f)" % (
+            self.shape, self.xc, self.yc, self.width,
+            self.height, self.rotation)
+
+
+class Regions:
+    """
+    Manipulate the region files (CIAO format), as well as parse regions
+    from strings.
+    """
+    REGION_SHAPES = {
+        "point": Point,
+        "circle": Circle,
+        "annulus": Annulus,
+        "pie": Pie,
+        "ellipse": Ellipse,
+        "box": Box,
+    }
+
+    def __init__(self, regfile=None):
+        if regfile:
+            self.regions = self.load(regfile)
+        else:
+            self.regions = []
+
+    def load(self, infile):
+        regstr = []
+        for line in open(infile):
+            if not (re.match(r"^\s*#.*$", line) or re.match(r"^\s*$", line)):
+                regstr.append(line.strip())
+        self.regions = self.parse(regstr)
+
+    def save(self, outfile, clobber=False):
+        if (not clobber) and os.path.exists(outfile):
+            raise OSError("output file already exists: %s" % outfile)
+        regstr = [str(reg) for reg in self.regions]
+        open(outfile, "w").write("\n".join(regstr) + "\n")
+
+    @classmethod
+    def parse(self, regstr):
+        """
+        Parse the given (list of) region string(s), and return the
+        corresponding parsed region objects.
+        """
+        if isinstance(regstr, list):
+            return [self.parse_single(reg) for reg in regstr]
+        else:
+            return self.parse_single(regstr)
+
+    @classmethod
+    def parse_single(self, regstr):
+        """
+        Parse the given single region string to its corresponding
+        region object.
+        """
+        shape = regstr.strip().split('(')[0].lower()
+        if shape in self.REGION_SHAPES:
+            return self.REGION_SHAPES[shape](regstr)
+        else:
+            raise ValueError("unknown region shape: %s" % shape)