Merge pull request #50 from bendudson/hdf5-coils

HDF5 improvements

freegs/coil.py

@@ -74,8 +74,8 @@ class Coil:
             (str("R"), np.float64),
             (str("Z"), np.float64),
             (str("current"), np.float64),
-            (str("turns"), np.int),
-            (str("control"), np.bool),
+            (str("turns"), int),
+            (str("control"), bool),
         ]
     )
 

freegs/dump.py

@@ -38,6 +38,8 @@
 
 from .equilibrium import Equilibrium
 from .machine import Coil, Circuit, Solenoid, Wall, Machine
+from .shaped_coil import ShapedCoil
+from .multi_coil import MultiCoil
 from . import boundary
 from . import machine
 
@@ -112,11 +114,15 @@ def write_equilibrium(self, equilibrium):
         """
 
         self.handle["coil_dtype"] = Coil.dtype
+        self.handle["shapedcoil_dtype"] = ShapedCoil.dtype
+        self.handle["multicoil_dtype"] = MultiCoil.dtype
         self.handle["circuit_dtype"] = Circuit.dtype
         self.handle["solenoid_dtype"] = Solenoid.dtype
 
         type_to_dtype = {
             Coil.dtype: self.handle["coil_dtype"],
+            ShapedCoil.dtype: self.handle["shapedcoil_dtype"],
+            MultiCoil.dtype: self.handle["multicoil_dtype"],
             Circuit.dtype: self.handle["circuit_dtype"],
             Solenoid.dtype: self.handle["solenoid_dtype"],
         }
@@ -164,7 +170,10 @@ def write_equilibrium(self, equilibrium):
 
         coils_group = tokamak_group.create_group(self.COILS_GROUP_NAME)
         for label, coil in equilibrium.tokamak.coils:
-            dtype = type_to_dtype[coil.dtype]
+            try:
+                dtype = type_to_dtype[coil.dtype]
+            except KeyError:
+                raise ValueError(f"Could not save coil type {type(coil)}")
             coils_group.create_dataset(
                 label, dtype=dtype, data=np.array(coil.to_numpy_array())
             )
@@ -186,6 +195,13 @@ def read_equilibrium(self):
         tokamak_group = equilibrium_group[self.MACHINE_GROUP_NAME]
         coil_group = tokamak_group[self.COILS_GROUP_NAME]
 
+        # HDF5 reads strings as bytes by default, so convert to string
+        def toString(s):
+            try:
+                return str(s, "utf-8")  # Convert bytes to string, using encoding
+            except TypeError:
+                return s  # Probably already a string
+
         # This is also a bit hacky - find the appropriate class in
         # freegs.machine and then call the `from_numpy_array` class
         # method
@@ -194,7 +210,9 @@ def make_coil_set(thing):
 
         # Unfortunately this creates the coils in lexographical order
         # by label, losing the origin
-        coils = [(label, make_coil_set(coil)) for label, coil in coil_group.items()]
+        coils = [
+            (toString(label), make_coil_set(coil)) for label, coil in coil_group.items()
+        ]
 
         if "wall_R" in tokamak_group:
             wall_R = tokamak_group["wall_R"][:]
@@ -218,6 +236,11 @@ def make_coil_set(thing):
         # string of the function __name__, which we then look up in
         # the boundary module dict
         eq_boundary_name = equilibrium_group["boundary_function"][()]
+
+        if hasattr(eq_boundary_name, "decode"):
+            # Convert bytes to string
+            eq_boundary_name = eq_boundary_name.decode()
+
         eq_boundary_func = boundary.__dict__[eq_boundary_name]
 
         equilibrium = Equilibrium(

freegs/machine.py

@@ -214,10 +214,21 @@ def from_numpy_array(cls, value):
                     this=type(cls), got=value.dtype, dtype=cls.dtype
                 )
             )
+
+        # HDF5 reads strings as bytes by default, so convert to string
+        def toString(s):
+            try:
+                return str(s, "utf-8")  # Convert bytes to string, using encoding
+            except TypeError:
+                return s  # Probably already a string
+
         # Use the current/control values from the first coil in the circuit
         # Should be consistent!
         return Circuit(
-            [(label, Coil(*coil), multiplier) for label, coil, multiplier in value],
+            [
+                (toString(label), Coil(*coil), multiplier)
+                for label, coil, multiplier in value
+            ],
             current=value[0][1]["current"] / value[0]["multiplier"],
             control=value[0][1]["control"],
         )
@@ -424,7 +435,7 @@ def __repr__(self):
         return "Wall(R={R}, Z={Z})".format(R=self.R, Z=self.Z)
 
     def __eq__(self, other):
-        return self.R == other.R and self.Z == other.Z
+        return np.allclose(self.R, other.R) and np.allclose(self.Z, other.Z)
 
     def __ne__(self, other):
         return not self == other

freegs/multi_coil.py

@@ -59,12 +59,14 @@ class MultiCoil(Coil):
     # A dtype for converting to Numpy array and storing in HDF5 files
     dtype = np.dtype(
         [
-            (str("R"), np.float64),
-            (str("Z"), np.float64),
+            (str("RZlen"), int),  # Length of R and Z arrays
+            (str("R"), "500f8"),  # Up to 100 points
+            (str("Z"), "500f8"),
             (str("current"), np.float64),
-            (str("turns"), np.int),
-            (str("control"), np.bool),
-            (str("mirror"), np.bool),
+            (str("turns"), int),
+            (str("control"), bool),
+            (str("mirror"), bool),
+            (str("polarity"), "2f8"),
         ]
     )
 
@@ -192,8 +194,26 @@ def to_numpy_array(self):
         """
         Helper method for writing output
         """
+
+        RZlen = len(self.Rfil)
+        assert RZlen <= 500
+        R = np.zeros(500)
+        Z = np.zeros(500)
+        R[:RZlen] = self.Rfil
+        Z[:RZlen] = self.Zfil
+
         return np.array(
-            (self.R, self.Z, self.current, self.turns, self.control), dtype=self.dtype
+            (
+                RZlen,
+                R,
+                Z,
+                self.current,
+                self.turns,
+                self.control,
+                self.mirror,
+                self.polarity,
+            ),
+            dtype=self.dtype,
         )
 
     @classmethod
@@ -204,7 +224,24 @@ def from_numpy_array(cls, value):
                     this=type(cls), got=value.dtype, dtype=cls.dtype
                 )
             )
-        return Coil(*value[()])
+        RZlen = value["RZlen"]
+        R = value["R"][:RZlen]
+        Z = value["Z"][:RZlen]
+        current = value["current"]
+        turns = value["turns"]
+        control = value["control"]
+        mirror = value["npoints"]
+        polarity = value["polarity"]
+
+        return MultiCoil(
+            R,
+            Z,
+            current=current,
+            turns=turns,
+            control=control,
+            mirror=mirror,
+            polarity=polarity,
+        )
 
     def plot(self, axis=None, show=False):
         """

freegs/shaped_coil.py

@@ -51,12 +51,13 @@ class ShapedCoil(Coil):
     # A dtype for converting to Numpy array and storing in HDF5 files
     dtype = np.dtype(
         [
-            (str("R"), np.float64),
-            (str("Z"), np.float64),
+            (str("RZlen"), int),  # Length of the R and Z arrays
+            (str("R"), "10f8"),  # Note: Up to 10 points
+            (str("Z"), "10f8"),  # Note: Up to 10 points
             (str("current"), np.float64),
-            (str("turns"), np.int),
-            (str("control"), np.bool),
-            (str("mirror"), np.bool),
+            (str("turns"), int),
+            (str("control"), bool),
+            (str("npoints"), int),
         ]
     )
 
@@ -86,6 +87,7 @@ def __init__(self, shape, current=0.0, turns=1, control=True, npoints=6):
         self.shape = shape
 
         # The quadrature points to be used
+        self.npoints_per_triangle = npoints
         self._points = quadrature.polygon_quad(shape, n=npoints)
 
     def controlPsi(self, R, Z):
@@ -176,3 +178,62 @@ def plot(self, axis=None, show=False):
         # axis.plot(rquad, zquad, 'ro')
 
         return axis
+
+    def to_numpy_array(self):
+        """
+        Helper method for writing output
+        """
+        RZlen = len(self.shape)
+        R = np.zeros(10)
+        Z = np.zeros(10)
+        R[:RZlen] = [R for R, Z in self.shape]
+        Z[:RZlen] = [Z for R, Z in self.shape]
+
+        return np.array(
+            (
+                RZlen,
+                R,
+                Z,
+                self.current,
+                self.turns,
+                self.control,
+                self.npoints_per_triangle,
+            ),
+            dtype=self.dtype,
+        )
+
+    @classmethod
+    def from_numpy_array(cls, value):
+        if value.dtype != cls.dtype:
+            raise ValueError(
+                "Can't create {this} from dtype: {got} (expected: {dtype})".format(
+                    this=type(cls), got=value.dtype, dtype=cls.dtype
+                )
+            )
+        RZlen = value["RZlen"]
+        R = value["R"][:RZlen]
+        Z = value["Z"][:RZlen]
+        current = value["current"]
+        turns = value["turns"]
+        control = value["control"]
+        npoints = value["npoints"]
+
+        return ShapedCoil(
+            list(zip(R, Z)),
+            current=current,
+            turns=turns,
+            control=control,
+            npoints=npoints,
+        )
+
+    def __eq__(self, other):
+        return (
+            np.allclose(self.shape, other.shape)
+            and self.current == other.current
+            and self.turns == other.turns
+            and self.control == other.control
+            and self.npoints_per_triangle == other.npoints_per_triangle
+        )
+
+    def __ne__(self, other):
+        return not self == other

freegs/test_readwrite.py

@@ -6,35 +6,37 @@
 
 def test_readwrite():
     """Test reading/writing to a file round-trip"""
-    tokamak = freegs.machine.MAST_sym()
-    eq = freegs.Equilibrium(
-        tokamak=tokamak,
-        Rmin=0.1,
-        Rmax=2.0,
-        Zmin=-1.0,
-        Zmax=1.0,
-        nx=17,
-        ny=17,
-        boundary=freegs.boundary.freeBoundaryHagenow,
-    )
-    profiles = freegs.jtor.ConstrainPaxisIp(1e4, 1e6, 2.0)
-
-    # Note here the X-point locations and isoflux locations are not the same.
-    # The result will be an unbalanced double null configuration, where the
-    # X-points are on different flux surfaces.
-    xpoints = [(1.1, -0.6), (1.1, 0.8)]
-    isoflux = [(1.1, -0.6, 1.1, 0.6)]
-    constrain = freegs.control.constrain(xpoints=xpoints, isoflux=isoflux)
-
-    freegs.solve(eq, profiles, constrain, maxits=25, atol=1e-3, rtol=1e-1)
-
-    memory_file = io.BytesIO()
-
-    with freegs.OutputFile(memory_file, "w") as f:
-        f.write_equilibrium(eq)
-
-    with freegs.OutputFile(memory_file, "r") as f:
-        read_eq = f.read_equilibrium()
-
-    assert tokamak == read_eq.tokamak
-    assert allclose(eq.psi(), read_eq.psi())
+
+    for tokamak in [freegs.machine.TestTokamak(), freegs.machine.MAST_sym()]:
+
+        eq = freegs.Equilibrium(
+            tokamak=tokamak,
+            Rmin=0.1,
+            Rmax=2.0,
+            Zmin=-1.0,
+            Zmax=1.0,
+            nx=17,
+            ny=17,
+            boundary=freegs.boundary.freeBoundaryHagenow,
+        )
+        profiles = freegs.jtor.ConstrainPaxisIp(1e4, 1e6, 2.0)
+
+        # Note here the X-point locations and isoflux locations are not the same.
+        # The result will be an unbalanced double null configuration, where the
+        # X-points are on different flux surfaces.
+        xpoints = [(1.1, -0.6), (1.1, 0.8)]
+        isoflux = [(1.1, -0.6, 1.1, 0.6)]
+        constrain = freegs.control.constrain(xpoints=xpoints, isoflux=isoflux)
+
+        freegs.solve(eq, profiles, constrain, maxits=25, atol=1e-3, rtol=1e-1)
+
+        memory_file = io.BytesIO()
+
+        with freegs.OutputFile(memory_file, "w") as f:
+            f.write_equilibrium(eq)
+
+        with freegs.OutputFile(memory_file, "r") as f:
+            read_eq = f.read_equilibrium()
+
+        assert tokamak == read_eq.tokamak
+        assert allclose(eq.psi(), read_eq.psi())

requirements.txt

@@ -1,4 +1,4 @@
 numpy>=1.14.1
 scipy>=1.0.0
 matplotlib>=2.0.0
-h5py==2.10.0
+h5py>=2.10.0

setup.py

@@ -22,7 +22,7 @@ def read(fname):
     install_requires=['numpy>=1.8',
                       'scipy>=0.14',
                       'matplotlib>=1.3',
-                      'h5py==2.10.0'],
+                      'h5py>=2.10.0'],
 
     platforms='any',