Support SpaceGroup-types in constrainAsSpaceGroup.

This closes #16.

Author: pavoljuhas

diffpy/srfit/structure/__init__.py

@@ -48,6 +48,6 @@ def struToParameterSet(name, stru):
 
     raise TypeError("Unadaptable structure format")
 
-from sgconstraints import constrainAsSpaceGroup
+from diffpy.srfit.structure.sgconstraints import constrainAsSpaceGroup
 
 # End of file

diffpy/srfit/structure/sgconstraints.py

@@ -22,7 +22,7 @@
 
 __all__ = ["constrainAsSpaceGroup"]
 
-def constrainAsSpaceGroup(phase, sgsymbol, scatterers = None,
+def constrainAsSpaceGroup(phase, spacegroup, scatterers = None,
         sgoffset = [0, 0, 0], constrainlat = True, constrainadps = True,
         adpsymbols = None, isosymbol = "Uiso"):
     """Constrain the structure to the space group.
@@ -33,8 +33,8 @@ def constrainAsSpaceGroup(phase, sgsymbol, scatterers = None,
 
     Arguments:
     phase       --  A BaseStructure object.
-    sgsymbol    --  The space group number or symbol (compatible with
-                    diffpy.Structure.SpaceGroups.GetSpaceGroup.
+    spacegroup  --  The space group number, symbol or an instance of
+                    SpaceGroup class from diffpy.Structure package.
     sgoffset    --  Optional offset for sg origin (default [0, 0, 0]).
     scatterers  --  The scatterer ParameterSets to constrain. If scatterers
                     is None (default), then all scatterers accessible from
@@ -77,9 +77,11 @@ def constrainAsSpaceGroup(phase, sgsymbol, scatterers = None,
 
     """
 
-    from diffpy.Structure.SpaceGroups import GetSpaceGroup
+    from diffpy.Structure.SpaceGroups import GetSpaceGroup, SpaceGroup
 
-    sg = GetSpaceGroup(sgsymbol)
+    sg = spacegroup
+    if not isinstance(spacegroup, SpaceGroup):
+        sg = GetSpaceGroup(spacegroup)
     sgp = _constrainAsSpaceGroup(phase, sg, scatterers, sgoffset,
             constrainlat, constrainadps, adpsymbols, isosymbol)
 

diffpy/srfit/tests/testsgconstriants.py

@@ -202,6 +202,24 @@ def _alltests(par):
 
         return
 
+
+    def test_ConstrainAsSpaceGroup_args(self):
+        """Test the arguments processing of constrainAsSpaceGroup function.
+        """
+        from diffpy.srfit.structure.sgconstraints import constrainAsSpaceGroup
+        from diffpy.Structure.SpaceGroups import GetSpaceGroup
+        stru = makeLaMnO3_P1()
+        parset = DiffpyStructureParSet("LaMnO3", stru)
+        sgpars = constrainAsSpaceGroup(parset, "P b n m")
+        sg = GetSpaceGroup('P b n m')
+        parset2 = DiffpyStructureParSet("LMO", makeLaMnO3_P1())
+        sgpars2 = constrainAsSpaceGroup(parset2, sg)
+        list(sgpars)
+        list(sgpars2)
+        self.assertEquals(sgpars.names, sgpars2.names)
+        return
+
+
 lamno3stru =\
 """\
 title  Cell structure file of LaMnO3.0